Merkle hash chip

src/circuit/gadget/sinsemilla.rs

@@ -9,6 +9,7 @@ use pasta_curves::arithmetic::{CurveAffine, FieldExt};
 use std::{convert::TryInto, fmt::Debug};
 
 pub mod chip;
+pub mod merkle;
 mod message;
 
 /// The set of circuit instructions required to use the [`Sinsemilla`](https://zcash.github.io/halo2/design/gadgets/sinsemilla.html) gadget.

src/circuit/gadget/sinsemilla/chip.rs

@@ -61,7 +61,7 @@ pub struct SinsemillaConfig {
     /// Fixed column shared by the whole circuit. This is used to load the
     /// x-coordinate of the domain $Q$, which is then constrained to equal the
     /// initial $x_a$.
-    constants: Column<Fixed>,
+    pub(super) constants: Column<Fixed>,
     /// Permutation over all advice columns and the `constants` fixed column.
     pub(super) perm: Permutation,
     /// Configure each advice column to be able to perform lookup range checks.
@@ -72,6 +72,13 @@ pub struct SinsemillaConfig {
     pub(super) lookup_config_4: LookupRangeCheckConfig<pallas::Base, { sinsemilla::K }>,
 }
 
+impl SinsemillaConfig {
+    /// Returns an array of all advice columns in this config, in arbitrary order.
+    pub(super) fn advices(&self) -> [Column<Advice>; 5] {
+        [self.x_a, self.x_p, self.bits, self.lambda_1, self.lambda_2]
+    }
+}
+
 #[derive(Eq, PartialEq, Clone, Debug)]
 pub struct SinsemillaChip {
     config: SinsemillaConfig,

src/circuit/gadget/sinsemilla/merkle.rs

@@ -0,0 +1,342 @@
+use halo2::{
+    circuit::{Chip, Layouter},
+    plonk::Error,
+};
+use pasta_curves::arithmetic::CurveAffine;
+
+use super::{HashDomains, SinsemillaInstructions};
+
+use crate::{
+    circuit::gadget::utilities::{
+        cond_swap::CondSwapInstructions, transpose_option_array, UtilitiesInstructions,
+    },
+    spec::i2lebsp,
+};
+use std::iter;
+
+mod chip;
+
+/// Instructions to check the validity of a Merkle path of a given `PATH_LENGTH`.
+/// The hash function used is a Sinsemilla instance with `K`-bit words.
+/// The hash function can process `MAX_WORDS` words.
+pub trait MerkleInstructions<
+    C: CurveAffine,
+    const PATH_LENGTH: usize,
+    const K: usize,
+    const MAX_WORDS: usize,
+>:
+    SinsemillaInstructions<C, K, MAX_WORDS>
+    + CondSwapInstructions<C::Base>
+    + UtilitiesInstructions<C::Base>
+    + Chip<C::Base>
+{
+    /// Compute MerkleCRH for a given `layer`. The hash that computes the root
+    /// is at layer 0, and the hashes that are applied to two leaves are at
+    /// layer `MERKLE_DEPTH_ORCHARD - 1` = layer 31.
+    #[allow(non_snake_case)]
+    fn hash_layer(
+        &self,
+        layouter: impl Layouter<C::Base>,
+        Q: C,
+        l: usize,
+        left: Self::Var,
+        right: Self::Var,
+    ) -> Result<Self::Var, Error>;
+}
+
+#[derive(Clone, Debug)]
+pub struct MerklePath<
+    C: CurveAffine,
+    MerkleChip,
+    const PATH_LENGTH: usize,
+    const K: usize,
+    const MAX_WORDS: usize,
+> where
+    MerkleChip: MerkleInstructions<C, PATH_LENGTH, K, MAX_WORDS> + Clone,
+{
+    chip_1: MerkleChip,
+    chip_2: MerkleChip,
+    domain: MerkleChip::HashDomains,
+    leaf_pos: Option<u32>,
+    // The Merkle path is ordered from leaves to root.
+    path: Option<[C::Base; PATH_LENGTH]>,
+}
+
+#[allow(non_snake_case)]
+impl<
+        C: CurveAffine,
+        MerkleChip,
+        const PATH_LENGTH: usize,
+        const K: usize,
+        const MAX_WORDS: usize,
+    > MerklePath<C, MerkleChip, PATH_LENGTH, K, MAX_WORDS>
+where
+    MerkleChip: MerkleInstructions<C, PATH_LENGTH, K, MAX_WORDS> + Clone,
+{
+    /// Calculates the root of the tree containing the given leaf at this Merkle path.
+    fn calculate_root(
+        &self,
+        mut layouter: impl Layouter<C::Base>,
+        leaf: MerkleChip::Var,
+    ) -> Result<MerkleChip::Var, Error> {
+        // A Sinsemilla chip uses 5 advice columns, but the full Orchard action circuit
+        // uses 10 advice columns. We distribute the path hashing across two Sinsemilla
+        // chips to make better use of the available circuit area.
+        let chips = iter::empty()
+            .chain(iter::repeat(self.chip_1.clone()).take(PATH_LENGTH / 2))
+            .chain(iter::repeat(self.chip_2.clone()));
+
+        // The Merkle path is ordered from leaves to root, which is consistent with the
+        // little-endian representation of `pos` below.
+        let path = transpose_option_array(self.path);
+
+        // Get position as a PATH_LENGTH-bit bitstring (little-endian bit order).
+        let pos: [Option<bool>; PATH_LENGTH] = {
+            let pos: Option<[bool; PATH_LENGTH]> = self.leaf_pos.map(|pos| i2lebsp(pos as u64));
+            transpose_option_array(pos)
+        };
+
+        let Q = self.domain.Q();
+
+        let mut node = leaf;
+        for (l, ((sibling, pos), chip)) in path.iter().zip(pos.iter()).zip(chips).enumerate() {
+            // `l` = 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` = 32 - 31 - 1 = 0.
+            // On the other hand, when `layer = 0` (the final sibling on the Merkle path),
+            // we have `l` = 32 - 0 - 1 = 31.
+            let pair = {
+                let pair = (node, *sibling);
+
+                // Swap node and sibling if needed
+                chip.swap(layouter.namespace(|| "node position"), pair, *pos)?
+            };
+
+            // Each `hash_layer` consists of 52 Sinsemilla words:
+            //  - l (10 bits) = 1 word
+            //  - left (255 bits) || right (255 bits) = 51 words (510 bits)
+            node = chip.hash_layer(
+                layouter.namespace(|| format!("hash l {}", l)),
+                Q,
+                l,
+                pair.0,
+                pair.1,
+            )?;
+        }
+
+        Ok(node)
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::{
+        chip::{MerkleChip, MerkleConfig},
+        MerklePath,
+    };
+
+    use crate::{
+        circuit::gadget::{
+            sinsemilla::chip::{SinsemillaChip, SinsemillaHashDomains},
+            utilities::{UtilitiesInstructions, Var},
+        },
+        constants::{L_ORCHARD_BASE, MERKLE_CRH_PERSONALIZATION, MERKLE_DEPTH_ORCHARD},
+        primitives::sinsemilla::HashDomain,
+        spec::i2lebsp,
+    };
+
+    use ff::PrimeFieldBits;
+    use halo2::{
+        arithmetic::FieldExt,
+        circuit::{layouter::SingleChipLayouter, Layouter},
+        dev::MockProver,
+        pasta::pallas,
+        plonk::{Assignment, Circuit, ConstraintSystem, Error},
+    };
+
+    use rand::random;
+    use std::convert::TryInto;
+
+    struct MyCircuit {
+        leaf: Option<pallas::Base>,
+        leaf_pos: Option<u32>,
+        merkle_path: Option<[pallas::Base; MERKLE_DEPTH_ORCHARD]>,
+    }
+
+    impl Circuit<pallas::Base> for MyCircuit {
+        type Config = (MerkleConfig, MerkleConfig);
+
+        fn configure(meta: &mut ConstraintSystem<pallas::Base>) -> Self::Config {
+            let advices = [
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+            ];
+
+            // Shared fixed column for loading constants
+            // TODO: Replace with public inputs API
+            let constants_1 = [
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+            ];
+            let constants_2 = [
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+            ];
+
+            let perm = meta.permutation(
+                &advices
+                    .iter()
+                    .map(|advice| (*advice).into())
+                    .chain(constants_1.iter().map(|fixed| (*fixed).into()))
+                    .chain(constants_2.iter().map(|fixed| (*fixed).into()))
+                    .collect::<Vec<_>>(),
+            );
+
+            // Fixed columns for the Sinsemilla generator lookup table
+            let lookup = (
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+            );
+
+            let sinsemilla_config_1 = SinsemillaChip::configure(
+                meta,
+                advices[5..].try_into().unwrap(),
+                lookup,
+                constants_1,
+                perm.clone(),
+            );
+            let config1 = MerkleChip::configure(meta, sinsemilla_config_1);
+
+            let sinsemilla_config_2 = SinsemillaChip::configure(
+                meta,
+                advices[..5].try_into().unwrap(),
+                lookup,
+                constants_2,
+                perm,
+            );
+            let config2 = MerkleChip::configure(meta, sinsemilla_config_2);
+
+            (config1, config2)
+        }
+
+        fn synthesize(
+            &self,
+            cs: &mut impl Assignment<pallas::Base>,
+            config: Self::Config,
+        ) -> Result<(), Error> {
+            let mut layouter = SingleChipLayouter::new(cs)?;
+
+            // Load generator table (shared across both configs)
+            SinsemillaChip::load(config.0.sinsemilla_config.clone(), &mut layouter)?;
+
+            // Construct Merkle chips which will be placed side-by-side in the circuit.
+            let chip_1 = MerkleChip::construct(config.0.clone());
+            let chip_2 = MerkleChip::construct(config.1.clone());
+
+            let leaf = chip_1.load_private(
+                layouter.namespace(|| ""),
+                config.0.cond_swap_config.a,
+                self.leaf,
+            )?;
+
+            let path = MerklePath {
+                chip_1,
+                chip_2,
+                domain: SinsemillaHashDomains::MerkleCrh,
+                leaf_pos: self.leaf_pos,
+                path: self.merkle_path,
+            };
+
+            let computed_final_root =
+                path.calculate_root(layouter.namespace(|| "calculate root"), leaf)?;
+
+            // The expected final root
+            let pos_bool = i2lebsp::<32>(self.leaf_pos.unwrap() as u64);
+            let path: Option<Vec<pallas::Base>> = self.merkle_path.map(|path| path.to_vec());
+            let final_root = hash_path(self.leaf.unwrap(), &pos_bool, &path.unwrap());
+
+            // Check the computed final root against the expected final root.
+            assert_eq!(computed_final_root.value().unwrap(), final_root);
+
+            Ok(())
+        }
+    }
+
+    fn hash_path(leaf: pallas::Base, pos_bool: &[bool], path: &[pallas::Base]) -> pallas::Base {
+        let domain = HashDomain::new(MERKLE_CRH_PERSONALIZATION);
+
+        // Compute the root
+        let mut node = leaf;
+        for (l, (sibling, pos)) in path.iter().zip(pos_bool.iter()).enumerate() {
+            let (left, right) = if *pos {
+                (*sibling, node)
+            } else {
+                (node, *sibling)
+            };
+
+            let l_star = i2lebsp::<10>(l as u64);
+            let left: Vec<_> = left
+                .to_le_bits()
+                .iter()
+                .by_val()
+                .take(L_ORCHARD_BASE)
+                .collect();
+            let right: Vec<_> = right
+                .to_le_bits()
+                .iter()
+                .by_val()
+                .take(L_ORCHARD_BASE)
+                .collect();
+
+            let mut message = l_star.to_vec();
+            message.extend_from_slice(&left);
+            message.extend_from_slice(&right);
+
+            node = domain.hash(message.into_iter()).unwrap();
+        }
+        node
+    }
+
+    #[test]
+    fn merkle_chip() {
+        // Choose a random leaf and position
+        let leaf = pallas::Base::rand();
+        let pos = random::<u32>();
+        let pos_bool = i2lebsp::<32>(pos as u64);
+
+        // Choose a path of random inner nodes
+        let path: Vec<_> = (0..(MERKLE_DEPTH_ORCHARD))
+            .map(|_| pallas::Base::rand())
+            .collect();
+
+        // This root is provided as a public input in the Orchard circuit.
+        let _root = hash_path(leaf, &pos_bool, &path);
+
+        let circuit = MyCircuit {
+            leaf: Some(leaf),
+            leaf_pos: Some(pos),
+            merkle_path: Some(path.try_into().unwrap()),
+        };
+
+        let prover = MockProver::run(11, &circuit, vec![]).unwrap();
+        assert_eq!(prover.verify(), Ok(()))
+    }
+}