[WIP] Implement MerkleInstructions

src/circuit/gadget/orchard_action/merkle.rs

@@ -5,60 +5,222 @@ use halo2::{
 use pasta_curves::arithmetic::{CurveAffine, FieldExt};
 use std::marker::PhantomData;
 
+use crate::constants::{util::i2lebsp, MERKLE_DEPTH_ORCHARD};
+
 use crate::circuit::gadget::{
     ecc::{
         chip::{EccChip, EccConfig},
         EccInstructions,
     },
-    sinsemilla::{SinsemillaChip, SinsemillaConfig, SinsemillaInstructions},
-    utilities::{UtilitiesChip, UtilitiesConfig, UtilitiesInstructions, Var},
+    sinsemilla::{HashDomains, SinsemillaChip, SinsemillaConfig, SinsemillaInstructions},
+    utilities::{Pair, Swap, UtilitiesChip, UtilitiesConfig, UtilitiesInstructions, Var},
 };
+use ff::PrimeField;
 use std::convert::TryInto;
 
-enum Node<F: FieldExt> {
-    Leaf(Var<F>),
-    Inner(Var<F>),
-    Sibling(Var<F>),
+#[derive(Clone, Debug)]
+pub struct MerkleConfig<const PATH_LENGTH: usize> {
+    utils_config: UtilitiesConfig,
+    sinsemilla_config: SinsemillaConfig,
 }
-struct Root<F: FieldExt>(pub Var<F>);
 
-pub trait MerkleInstructions<F: FieldExt, const MERKLE_DEPTH: usize>:
-    UtilitiesInstructions<F>
+pub trait MerkleInstructions<C: CurveAffine, const PATH_LENGTH: usize>:
+    SinsemillaInstructions<C, X = Var<C::Base>>
+    + UtilitiesInstructions<C::Base, Var = Var<C::Base>, Pair = Pair<C::Base>, Swap = Swap>
 {
     /// Check the validity of a Merkle path from a given leaf to a claimed root.
+    #[allow(non_snake_case)]
     fn merkle_path_check(
         &self,
-        layouter: impl Layouter<F>,
-        root: Option<[u8; 32]>,
-        leaf: Option<[u8; 32]>,
-        merkle_path: [Option<[u8; 32]>; MERKLE_DEPTH],
+        mut layouter: impl Layouter<C::Base>,
+        domain: &<Self as SinsemillaInstructions<C>>::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);
+
+        // 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(|| "hash leaf"), leaf.0)?;
+        for (layer, (sibling, pos)) in merkle_path.iter().zip(pos.iter()).enumerate() {
+            // Load sibling into circuit
+            let sibling = self.load_private(
+                layouter.namespace(|| format!("sibling {}", layer)),
+                *sibling,
+            )?;
+            let pair: Pair<C::Base> = (node, sibling).into();
+
+            // Swap node and sibling if needed
+            let swap = pos.map(|pos| C::Base::from_u64(pos as u64));
+            let swap: Swap = self
+                .load_private(layouter.namespace(|| format!("pos {}", layer)), swap)?
+                .into();
+
+            let (left, right) = layouter.assign_region(
+                || "swap",
+                |mut region| {
+                    let pair: Pair<C::Base> = self.swap(&mut region, layer, pair, swap)?;
+                    let left = pair.x();
+                    let right = pair.y();
+
+                    Ok((left, right))
+                },
+            )?;
+
+            node = self.layer_hash(
+                layouter.namespace(|| format!("hash layer {}", layer)),
+                Q,
+                layer,
+                left,
+                right,
+            )?;
+        }
+
         Ok(())
     }
+
+    fn layer_hash(
+        &self,
+        mut layouter: impl Layouter<C::Base>,
+        Q: C,
+        layer: usize,
+        left: Var<C::Base>,
+        right: Var<C::Base>,
+    ) -> Result<Var<C::Base>, Error> {
+        // We need to hash `l || left || right`, where `l` is a 10-bit value.
+        // `a = a_0||a_1` = `l` || (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 = {
+            // a_0 = layer
+            let a_0 = (C::Base::from_u64(layer 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.map(|a_1| prepare_message_piece(&[a_0.into(), a_1.into()], 250));
+            let a = transpose_option_vec(a, 250);
+
+            self.witness_message_piece_bitstring(layouter.namespace(|| "a"), a, 25)?
+        };
+
+        // `b = b_0||b_1` = (bits 240..254 of `left`) || (bits 0..234 of `right`)
+        let 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
+                .zip(b_1)
+                .map(|(b_0, b_1)| prepare_message_piece(&[b_0.into(), b_1.into()], 250));
+            let b = transpose_option_vec(b, 250);
+
+            self.witness_message_piece_bitstring(layouter.namespace(|| "b"), b, 25)?
+        };
+
+        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()], 20));
+            let c = transpose_option_vec(c, 20);
+            self.witness_message_piece_bitstring(layouter.namespace(|| "c"), c, 2)?
+        };
+
+        let point = self.hash_to_point(
+            layouter.namespace(|| format!("layer {}", layer)),
+            Q,
+            vec![a, b, c],
+        )?;
+        Ok(Self::extract_p(&point).clone())
+    }
 }
 
-fn layer_hash<C: CurveAffine>(
-    chip: MerkleChip<C>,
-    layer: u16,
-    left: Node<C::Base>,
-    right: Node<C::Base>,
-) -> Result<Var<C::Base>, Error> {
-    todo!()
+struct MessageSubPiece<F: FieldExt> {
+    field_elem: F,
+    bit_range: std::ops::Range<usize>,
 }
 
-#[derive(Clone, Debug)]
-pub struct MerkleConfig {
-    config1: (UtilitiesConfig, SinsemillaConfig),
-    config2: (UtilitiesConfig, SinsemillaConfig),
+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 Vec<Option<T>> given an Option<Vec<T>> of a certain length.
+fn transpose_option_vec<T: Copy + Clone>(
+    option_vec: Option<Vec<T>>,
+    length: usize,
+) -> Vec<Option<T>> {
+    if let Some(vec) = option_vec {
+        assert_eq!(length, vec.len());
+        vec.iter().map(|elem| Some(*elem)).collect()
+    } else {
+        vec![None; length]
+    }
+}
+
+// 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
+    })
 }
 
-pub struct MerkleChip<C: CurveAffine> {
-    config: MerkleConfig,
+pub struct MerkleChip<C: CurveAffine, const PATH_LENGTH: usize> {
+    config: MerkleConfig<PATH_LENGTH>,
     _marker: PhantomData<C>,
 }
 
-impl<C: CurveAffine> Chip<C::Base> for MerkleChip<C> {
-    type Config = MerkleConfig;
+impl<C: CurveAffine, const PATH_LENGTH: usize> Chip<C::Base> for MerkleChip<C, PATH_LENGTH> {
+    type Config = MerkleConfig<PATH_LENGTH>;
     type Loaded = ();
 
     fn config(&self) -> &Self::Config {
@@ -70,46 +232,60 @@ impl<C: CurveAffine> Chip<C::Base> for MerkleChip<C> {
     }
 }
 
-impl<C: CurveAffine> MerkleChip<C> {
+impl<C: CurveAffine, const PATH_LENGTH: usize> MerkleChip<C, PATH_LENGTH> {
     pub fn configure(
         meta: &mut ConstraintSystem<C::Base>,
         advices: [Column<Advice>; 10],
         perm: Permutation,
-    ) -> MerkleConfig {
+    ) -> MerkleConfig<PATH_LENGTH> {
         let ecc_config = EccChip::<C>::configure(meta, advices);
 
         let lookup = (
             meta.fixed_column(),
             meta.fixed_column(),
             meta.fixed_column(),
         );
-        let config1 = (
-            UtilitiesChip::configure(meta, advices.clone()[..5].try_into().unwrap(), perm.clone()),
-            SinsemillaChip::<C>::configure(
-                meta,
-                ecc_config.clone(),
-                advices.clone()[..5].try_into().unwrap(),
-                lookup,
-                perm.clone(),
-            ),
-        );
-        let config2 = (
-            UtilitiesChip::configure(meta, advices.clone()[5..].try_into().unwrap(), perm.clone()),
-            SinsemillaChip::<C>::configure(
-                meta,
-                ecc_config,
-                advices.clone()[5..].try_into().unwrap(),
-                lookup,
-                perm.clone(),
-            ),
+        let utils_config =
+            UtilitiesChip::configure(meta, advices.clone()[..5].try_into().unwrap(), perm.clone());
+        let sinsemilla_config = SinsemillaChip::<C>::configure(
+            meta,
+            ecc_config.clone(),
+            advices.clone()[..5].try_into().unwrap(),
+            lookup,
+            perm.clone(),
         );
-        MerkleConfig { config1, config2 }
+
+        MerkleConfig {
+            utils_config,
+            sinsemilla_config,
+        }
     }
 
-    pub fn construct(config: MerkleConfig) -> Self {
+    pub fn construct(config: MerkleConfig<PATH_LENGTH>) -> Self {
         MerkleChip {
             config,
             _marker: PhantomData,
         }
     }
 }
+
+// impl<C: CurveAffine, const PATH_LENGTH: usize> MerkleInstructions<C, PATH_LENGTH>
+//     for MerkleChip<C, PATH_LENGTH>
+// {
+// }
+
+// impl<C: CurveAffine, const PATH_LENGTH: usize> UtilitiesInstructions<C::Base>
+//     for MerkleChip<C, PATH_LENGTH>
+// {
+// }
+
+// impl<C: CurveAffine, const PATH_LENGTH: usize> SinsemillaInstructions<C>
+//     for MerkleChip<C, PATH_LENGTH>
+// {
+// }
+
+#[cfg(test)]
+pub mod tests {
+    #[test]
+    fn merkle_path_check() {}
+}

src/circuit/gadget/utilities.rs

@@ -8,14 +8,15 @@ use std::marker::PhantomData;
 pub mod cond_swap;
 pub mod plonk;
 
-use cond_swap::{CondSwapChip, CondSwapConfig, CondSwapInstructions, Pair, Swap};
+use cond_swap::{CondSwapChip, CondSwapConfig, CondSwapInstructions};
+pub use cond_swap::{Pair, Swap};
 use plonk::{PLONKChip, PLONKConfig, PLONKInstructions};
 
 /// A variable representing a number.
 #[derive(Copy, Clone)]
 pub struct Var<F: FieldExt> {
-    cell: Cell,
-    value: Option<F>,
+    pub cell: Cell,
+    pub value: Option<F>,
 }
 
 #[derive(Clone, Debug)]

src/constants/util.rs

@@ -1,5 +1,6 @@
 use ff::PrimeField;
 use halo2::arithmetic::{CurveAffine, FieldExt};
+use std::convert::TryInto;
 
 /// Decompose a scalar into `window_num_bits` bits (little-endian)
 /// For a window size of `w`, this returns [k_0, ..., k_n] where each `k_i`
@@ -40,6 +41,21 @@ pub fn evaluate<C: CurveAffine>(x: u8, coeffs: &[C::Base]) -> C::Base {
         .fold(C::Base::default(), |acc, coeff| acc * x + coeff)
 }
 
+/// The sequence of bits representing a u32 in little-endian order.
+pub fn i2lebsp<const PATH_LENGTH: usize>(int: u32) -> [bool; PATH_LENGTH] {
+    let mut int = int;
+    (0..PATH_LENGTH)
+        .rev()
+        .map(|shift| {
+            let bit = int >> shift;
+            int = int >> 1;
+            bit == 1
+        })
+        .collect::<Vec<_>>()
+        .try_into()
+        .unwrap()
+}
+
 #[cfg(test)]
 mod tests {
     use super::decompose_scalar_fixed;