Implement MerkleInstructions for MerkleChip

src/circuit/gadget/orchard_action/merkle.rs

@@ -180,6 +180,379 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
     }
 }
 
+struct MessageSubPiece<F: FieldExt> {
+    field_elem: F,
+    bit_range: std::ops::Range<usize>,
+}
+
+impl<F: FieldExt> From<(F, std::ops::Range<usize>)> for MessageSubPiece<F> {
+    fn from(field_elem_bit_range: (F, std::ops::Range<usize>)) -> Self {
+        Self {
+            field_elem: field_elem_bit_range.0,
+            bit_range: field_elem_bit_range.1,
+        }
+    }
+}
+
+// Returns a message piece given a slice of `MessageSubPiece`s, which are field
+// elements with specified bit-ranges to be used in the message piece.
+fn prepare_message_piece<F: FieldExt>(
+    subpieces: &[MessageSubPiece<F>],
+    num_words: usize,
+) -> Vec<bool> {
+    // The number of words in the message piece should be consistent with the number of bits used.
+    assert_eq!(
+        num_words * 10,
+        subpieces
+            .iter()
+            .map(|subpiece| subpiece.bit_range.len())
+            .sum()
+    );
+
+    subpieces.iter().fold(Vec::new(), |mut acc, subpiece| {
+        let bits: Vec<_> = subpiece
+            .field_elem
+            .to_le_bits()
+            .iter()
+            .by_val()
+            .take(F::NUM_BITS as usize)
+            .collect();
+        let bits = bits[subpiece.bit_range.clone()].to_vec();
+        acc.extend_from_slice(&bits);
+        acc
+    })
+}
+
+
+impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS: usize>
+    MerkleInstructions<C, PATH_LENGTH, K, MAX_WORDS> for MerkleChip<C, PATH_LENGTH, K, MAX_WORDS>
+{
+    /// Check the validity of a Merkle path from a given leaf to a claimed root.
+    #[allow(non_snake_case)]
+    fn merkle_path_check(
+        &self,
+        mut layouter: impl Layouter<C::Base>,
+        domain: &<Self as SinsemillaInstructions<C, K, MAX_WORDS>>::HashDomains,
+        root: Option<C::Base>,
+        leaf: (Option<C::Base>, Option<u32>),
+        merkle_path: Vec<Option<C::Base>>,
+    ) -> Result<(), Error> {
+        assert_eq!(merkle_path.len(), PATH_LENGTH);
+
+        let config = self.config().clone();
+
+        // Get position as a 32-bit bitstring (little-endian bit order).
+        let pos = leaf.1;
+        let pos: Option<[bool; PATH_LENGTH]> = pos.map(i2lebsp);
+        let pos: [Option<bool>; PATH_LENGTH] = if let Some(pos) = pos {
+            pos.iter()
+                .map(|pos| Some(*pos))
+                .collect::<Vec<_>>()
+                .try_into()
+                .unwrap()
+        } else {
+            [None; PATH_LENGTH]
+        };
+
+        let Q = domain.Q();
+
+        let mut node =
+            self.load_private(layouter.namespace(|| ""), config.cond_swap_config.x, leaf.0)?;
+
+        for (l_star, (sibling, pos)) in merkle_path.iter().zip(pos.iter()).enumerate() {
+            // `l_star` = MERKLE_DEPTH_ORCHARD - layer - 1, which is the index obtained from
+            // enumerating this Merkle path (going from leaf to root).
+            // For example, when `layer = 31` (the first sibling on the Merkle path),
+            // we have `l_star` = 32 - 31 - 1 = 0.
+            // On the other hand, when `layer = 0` (the final sibling on the Merkle path),
+            // we have `l_star` = 32 - 0 - 1 = 31.
+
+            let pair = {
+                // Load sibling into circuit
+                let sibling = self.load_private(
+                    layouter.namespace(|| ""),
+                    config.cond_swap_config.x,
+                    *sibling,
+                )?;
+                let pair = (node, sibling);
+
+                // Swap node and sibling if needed
+                let swap = pos.map(|pos| C::Base::from_u64(pos as u64));
+                let swap = self
+                    .load_private(
+                        layouter.namespace(|| format!("pos {}", l_star)),
+                        config.cond_swap_config.swap,
+                        swap,
+                    )?
+                    .into();
+
+                self.swap(layouter.namespace(|| ""), pair, swap)?
+            };
+
+            // Each `layer_hash` consists of 52 Sinsemilla words.
+            node = self.layer_hash(
+                layouter.namespace(|| format!("hash l_star {}", l_star)),
+                Q,
+                l_star,
+                pair.0,
+                pair.1,
+            )?;
+        }
+
+        // Check that the final node is equal to the claimed root.
+        layouter.assign_region(
+            || "Root check",
+            |mut region| {
+                region.assign_advice(
+                    || "Witness root",
+                    config.a,
+                    0,
+                    || root.ok_or(Error::SynthesisError),
+                )?;
+
+                // Copy final output of Sinsemilla hash
+                copy(
+                    &mut region,
+                    || "Copy final hash output",
+                    config.b,
+                    0,
+                    &node,
+                    &config.perm,
+                )?;
+
+                config.q_merkle.enable(&mut region, 0)
+            },
+        )
+    }
+
+    #[allow(non_snake_case)]
+    fn layer_hash(
+        &self,
+        mut layouter: impl Layouter<C::Base>,
+        Q: C,
+        l_star: usize,
+        left: CellValue<C::Base>,
+        right: CellValue<C::Base>,
+    ) -> Result<CellValue<C::Base>, Error> {
+        // <https://zips.z.cash/protocol/nu5.pdf#orchardmerklecrh>
+        // We need to hash `l_star || left || right`, where `l_star` is a 10-bit value.
+        // We allow `left` and `right` to be non-canonical 255-bit encodings.
+        //
+        // `a = a_0||a_1` = `l_star` || (bits 0..=239 of `left`)
+        // `b = b_0||b_1` = (bits 240..=254 of `left`) || (bits 0..=234 of `right`)
+        // `c = bits 235..=254 of `right`
+
+        // `a = a_0||a_1` = `l` || (bits 0..=239 of `left`)
+        let (a_0, a_1, a) = {
+            // a_0 = l_star
+            let a_0 = (C::Base::from_u64(l_star as u64), 0..10);
+
+            // a_1 = (bits 0..=239 of `left`)
+            let a_1 = left.value().map(|left| (left, 0..240));
+
+            let a: Option<Vec<bool>> = a_1
+                .clone()
+                .map(|a_1| prepare_message_piece(&[a_0.clone().into(), a_1.into()], 25));
+
+            let a = transpose_option_vec(a, 250);
+            let a = self.witness_message_piece_bitstring(layouter.namespace(|| "a"), a, 25)?;
+
+            let a_1 = a_1.map(|a_1| {
+                to_field_elem(
+                    &a_1.0
+                        .to_le_bits()
+                        .iter()
+                        .by_val()
+                        .take(240)
+                        .collect::<Vec<_>>(),
+                )
+            });
+
+            (a_0.0, a_1, a)
+        };
+
+        // `b = b_0||b_1` = (bits 240..=254 of `left`) || (bits 0..=234 of `right`)
+        let (b_0, b_1, b) = {
+            // b_0 = (bits 240..=254 of `left`)
+            let b_0 = left
+                .value()
+                .map(|left| (left, 240..(C::Base::NUM_BITS as usize)));
+
+            // b_1 = (bits 0..=234 of `right`)
+            let b_1 = right.value().map(|right| (right, 0..235));
+
+            let b: Option<Vec<bool>> = b_0
+                .clone()
+                .zip(b_1.clone())
+                .map(|(b_0, b_1)| prepare_message_piece(&[b_0.into(), b_1.into()], 25));
+
+            let b = transpose_option_vec(b, 250);
+            let b = self.witness_message_piece_bitstring(layouter.namespace(|| "b"), b, 25)?;
+
+            let b_0 = b_0.map(|b_0| {
+                let b_0 = &b_0
+                    .0
+                    .to_le_bits()
+                    .iter()
+                    .by_val()
+                    .take(C::Base::NUM_BITS as usize)
+                    .collect::<Vec<_>>()[b_0.1];
+                to_field_elem(b_0)
+            });
+
+            let b_1 = b_1.map(|b_1| {
+                let b_1 = &b_1
+                    .0
+                    .to_le_bits()
+                    .iter()
+                    .by_val()
+                    .take(C::Base::NUM_BITS as usize)
+                    .collect::<Vec<_>>()[b_1.1];
+                to_field_elem(b_1)
+            });
+
+            (b_0, b_1, b)
+        };
+
+        let c = {
+            // `c = bits 235..=254 of `right`
+            let c = right
+                .value()
+                .map(|right| (right, 235..(C::Base::NUM_BITS as usize)));
+            let c: Option<Vec<bool>> = c.map(|c| prepare_message_piece(&[c.into()], 2));
+
+            let c = transpose_option_vec(c, 20);
+            self.witness_message_piece_bitstring(layouter.namespace(|| "c"), c, 2)?
+        };
+
+        let config = self.config().clone();
+        // Check that the pieces have been decomposed properly.
+        {
+            layouter.assign_region(
+                || "Check piece decomposition",
+                |mut region| {
+                    // Set the fixed column `l_star_plus1` to the current l_star + 1.
+                    let l_star_plus1 = (l_star as u64) + 1;
+                    region.assign_fixed(
+                        || format!("l_star_plus1 {}", l_star_plus1),
+                        config.l_star,
+                        0,
+                        || Ok(C::Base::from_u64(l_star_plus1)),
+                    )?;
+
+                    // Copy and assign `a` at the correct position.
+                    copy(
+                        &mut region,
+                        || "copy a",
+                        config.a,
+                        0,
+                        &a.cell(),
+                        &config.perm,
+                    )?;
+                    // Copy and assign `b` at the correct position.
+                    copy(
+                        &mut region,
+                        || "copy b",
+                        config.b,
+                        0,
+                        &b.cell(),
+                        &config.perm,
+                    )?;
+                    // Copy and assign `c` at the correct position.
+                    copy(
+                        &mut region,
+                        || "copy c",
+                        config.c,
+                        0,
+                        &c.cell(),
+                        &config.perm,
+                    )?;
+                    // Copy and assign the left node at the correct position.
+                    copy(&mut region, || "left", config.left, 0, &left, &config.perm)?;
+                    // Copy and assign the right node at the correct position.
+                    copy(
+                        &mut region,
+                        || "right",
+                        config.right,
+                        0,
+                        &right,
+                        &config.perm,
+                    )?;
+
+                    // Witness the piece `a_0`.
+                    region.assign_advice(|| "a_0", config.a, 1, || Ok(a_0))?;
+                    // Witness the piece `a_1`.
+                    region.assign_advice(
+                        || "a_1",
+                        config.b,
+                        1,
+                        || a_1.ok_or(Error::SynthesisError),
+                    )?;
+                    // Witness the piece `b_0`.
+                    region.assign_advice(
+                        || "b_0",
+                        config.c,
+                        1,
+                        || b_0.ok_or(Error::SynthesisError),
+                    )?;
+                    // Witness the piece `b_1`.
+                    region.assign_advice(
+                        || "b_1",
+                        config.left,
+                        1,
+                        || b_1.ok_or(Error::SynthesisError),
+                    )?;
+
+                    Ok(())
+                },
+            )?;
+        }
+
+        let point = self.hash_to_point(
+            layouter.namespace(|| format!("l_star {}", l_star)),
+            Q,
+            vec![a, b, c],
+        )?;
+        let result = *Self::extract(&point);
+
+        // Check layer hash output against Sinsemilla primitives hash
+        #[cfg(test)]
+        {
+            use crate::constants::{L_ORCHARD_BASE, MERKLE_CRH_PERSONALIZATION};
+            use crate::primitives::sinsemilla::HashDomain;
+
+            let l_star = i2lebsp::<10>(l_star as u32);
+            let left: Vec<_> = left
+                .value()
+                .unwrap()
+                .to_le_bits()
+                .iter()
+                .by_val()
+                .take(L_ORCHARD_BASE)
+                .collect();
+            let right: Vec<_> = right
+                .value()
+                .unwrap()
+                .to_le_bits()
+                .iter()
+                .by_val()
+                .take(L_ORCHARD_BASE)
+                .collect();
+            let merkle_crh = HashDomain::new(MERKLE_CRH_PERSONALIZATION);
+
+            let mut message = l_star.to_vec();
+            message.extend_from_slice(&left);
+            message.extend_from_slice(&right);
+
+            let expected = merkle_crh.hash(message.into_iter()).unwrap();
+
+            assert_eq!(expected.to_bytes(), result.value().unwrap().to_bytes());
+        }
+
+        Ok(result)
+    }
+}
+
 impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS: usize> Chip<C::Base>
     for MerkleChip<C, PATH_LENGTH, K, MAX_WORDS>
 {