1use crate::crh::FixedLengthCRH;
  2use algebra::bytes::ToBytes;
  3use crate::Error;
  4use std::{fmt, rc::Rc};
  5
  6
  7#[cfg(feature = "r1cs")]
  8pub mod constraints;
  9
 10pub trait MerkleTreeConfig {
 11    const HEIGHT: usize;
 12    type H: FixedLengthCRH;
 13}
 14
 15/// Stores the hashes of a particular path (in order) from leaf to root.
 16/// Our path `is_left_child()` if the boolean in `path` is true.
 17#[derive(Derivative)]
 18#[derivative(
 19    Clone(bound = "P: MerkleTreeConfig"),
 20    Debug(bound = "P: MerkleTreeConfig, <P::H as FixedLengthCRH>::Output: fmt::Debug")
 21)]
 22pub struct MerkleTreePath<P: MerkleTreeConfig> {
 23    pub(crate) path: Vec<(<P::H as FixedLengthCRH>::Output, <P::H as FixedLengthCRH>::Output)>,
 24}
 25
 26pub type MerkleTreeParams<P> = <<P as MerkleTreeConfig>::H as FixedLengthCRH>::Parameters;
 27pub type MerkleTreeDigest<P> = <<P as MerkleTreeConfig>::H as FixedLengthCRH>::Output;
 28
 29
 30impl<P: MerkleTreeConfig> Default for MerkleTreePath<P> {
 31    fn default() -> Self {
 32        let mut path = Vec::with_capacity(P::HEIGHT as usize);
 33        for _i in 1..P::HEIGHT as usize {
 34            path.push((<P::H as FixedLengthCRH>::Output::default(), <P::H as FixedLengthCRH>::Output::default()));
 35        }
 36        Self {
 37            path,
 38        }
 39    }
 40}
 41
 42impl<P: MerkleTreeConfig> MerkleTreePath<P> {
 43    pub fn verify<L: ToBytes>(
 44        &self,
 45        parameters: &<P::H as FixedLengthCRH>::Parameters,
 46        root_hash: &<P::H as FixedLengthCRH>::Output,
 47        leaf: &L,
 48    ) -> Result<bool, Error> {
 49        if self.path.len() != (P::HEIGHT - 1) as usize {
 50            return Ok(false);
 51        }
 52        // Check that the given leaf matches the leaf in the membership proof.
 53        let mut buffer = [0u8; 128];
 54
 55        if !self.path.is_empty() {
 56            let claimed_leaf_hash = hash_leaf::<P::H, L>(parameters, leaf, &mut buffer)?;
 57
 58            // Check if leaf is one of the bottom-most siblings.
 59            if claimed_leaf_hash != self.path[0].0 && claimed_leaf_hash != self.path[0].1 {
 60                return Ok(false);
 61            };
 62
 63            let mut prev = claimed_leaf_hash;
 64            // Check levels between leaf level and root.
 65            for &(ref hash, ref sibling_hash) in &self.path {
 66                // Check if the previous hash matches the correct current hash.
 67                if &prev != hash && &prev != sibling_hash {
 68                    return Ok(false);
 69                };
 70                prev = hash_inner_node::<P::H>(parameters, hash, sibling_hash, &mut buffer)?;
 71            }
 72
 73            if root_hash != &prev {
 74                return Ok(false);
 75            }
 76            Ok(true)
 77        } else {
 78            Ok(false)
 79        }
 80    }
 81}
 82
 83pub struct MerkleHashTree<P: MerkleTreeConfig> {
 84    tree:         Vec<<P::H as FixedLengthCRH>::Output>,
 85    padding_tree: Vec<(<P::H as FixedLengthCRH>::Output, <P::H as FixedLengthCRH>::Output)>,
 86    parameters:   Rc<<P::H as FixedLengthCRH>::Parameters>,
 87    root:         Option<<P::H as FixedLengthCRH>::Output>,
 88}
 89
 90impl<P: MerkleTreeConfig> MerkleHashTree<P> {
 91    pub const HEIGHT: u8 = P::HEIGHT as u8;
 92
 93    pub fn blank(parameters: Rc<<P::H as FixedLengthCRH>::Parameters>) -> Self {
 94        MerkleHashTree {
 95            tree: Vec::new(),
 96            padding_tree: Vec::new(),
 97            root: None,
 98            parameters,
 99        }
100    }
101
102    pub fn new<L: ToBytes>(parameters: Rc<<P::H as FixedLengthCRH>::Parameters>, leaves: &[L]) -> Result<Self, Error> {
103        let new_time = start_timer!(|| "MerkleTree::New");
104
105        let last_level_size = leaves.len().next_power_of_two();
106        let tree_size = 2 * last_level_size - 1;
107        let tree_height = tree_height(tree_size);
108        assert!(tree_height as u8 <= Self::HEIGHT);
109
110        // Initialize the merkle tree.
111        let mut tree = Vec::with_capacity(tree_size);
112        let empty_hash = hash_empty::<P::H>(&parameters)?;
113        for _ in 0..tree_size {
114            tree.push(empty_hash.clone());
115        }
116
117        // Compute the starting indices for each level of the tree.
118        let mut index = 0;
119        let mut level_indices = Vec::with_capacity(tree_height);
120        for _ in 0..tree_height {
121            level_indices.push(index);
122            index = left_child(index);
123        }
124
125        // Compute and store the hash values for each leaf.
126        let last_level_index = level_indices.pop().unwrap();
127        let mut buffer = [0u8; 128];
128        for (i, leaf) in leaves.iter().enumerate() {
129            tree[last_level_index + i] = hash_leaf::<P::H, _>(&parameters, leaf, &mut buffer)?;
130        }
131
132        // Compute the hash values for every node in the tree.
133        let mut upper_bound = last_level_index;
134        let mut buffer = [0u8; 128];
135        level_indices.reverse();
136        for &start_index in &level_indices {
137            // Iterate over the current level.
138            for current_index in start_index..upper_bound {
139                let left_index = left_child(current_index);
140                let right_index = right_child(current_index);
141
142                // Compute Hash(left || right).
143                tree[current_index] = hash_inner_node::<P::H>(
144                    &parameters,
145                    &tree[left_index],
146                    &tree[right_index],
147                    &mut buffer,
148                )?;
149            }
150            upper_bound = start_index;
151        }
152        // Finished computing actual tree.
153        // Now, we compute the dummy nodes until we hit our HEIGHT goal.
154        let mut cur_height = tree_height;
155        let mut padding_tree = Vec::new();
156        let mut cur_hash = tree[0].clone();
157        while cur_height < (Self::HEIGHT - 1) as usize {
158            cur_hash = hash_inner_node::<P::H>(&parameters, &cur_hash, &empty_hash, &mut buffer)?;
159            padding_tree.push((cur_hash.clone(), empty_hash.clone()));
160            cur_height += 1;
161        }
162
163        let root_hash = hash_inner_node::<P::H>(&parameters, &cur_hash, &empty_hash, &mut buffer)?;
164
165        end_timer!(new_time);
166
167        Ok(MerkleHashTree {
168            tree,
169            padding_tree,
170            parameters,
171            root: Some(root_hash),
172        })
173    }
174
175    #[inline]
176    pub fn root(&self) -> <P::H as FixedLengthCRH>::Output {
177        self.root.clone().unwrap()
178    }
179
180    pub fn generate_proof<L: ToBytes>(
181        &self,
182        index: usize,
183        leaf: &L,
184    ) -> Result<MerkleTreePath<P>, Error> {
185        let prove_time = start_timer!(|| "MerkleTree::GenProof");
186        let mut path = Vec::new();
187
188        let mut buffer = [0u8; 128];
189        let leaf_hash = hash_leaf::<P::H, _>(&self.parameters, leaf, &mut buffer)?;
190        let tree_height = tree_height(self.tree.len());
191        let tree_index = convert_index_to_last_level(index, tree_height);
192        let empty_hash = hash_empty::<P::H>(&self.parameters)?;
193
194        // Check that the given index corresponds to the correct leaf.
195        if leaf_hash != self.tree[tree_index] {
196            Err(MerkleTreeError::IncorrectLeafIndex(tree_index))?
197        }
198
199        // Iterate from the leaf up to the root, storing all intermediate hash values.
200        let mut current_node = tree_index;
201        while !is_root(current_node) {
202            let sibling_node = sibling(current_node).unwrap();
203            let (curr_hash, sibling_hash) = (
204                self.tree[current_node].clone(),
205                self.tree[sibling_node].clone(),
206            );
207            if is_left_child(current_node) {
208                path.push((curr_hash, sibling_hash));
209            } else {
210                path.push((sibling_hash, curr_hash));
211            }
212            current_node = parent(current_node).unwrap();
213        }
214
215        // Store the root node. Set boolean as true for consistency with digest
216        // location.
217        assert!(path.len() < Self::HEIGHT as usize);
218        if path.len() != (Self::HEIGHT - 1) as usize {
219            path.push((self.tree[0].clone(), empty_hash));
220            for &(ref hash, ref sibling_hash) in &self.padding_tree {
221                path.push((hash.clone(), sibling_hash.clone()));
222            }
223        }
224        end_timer!(prove_time);
225        if path.len() != (Self::HEIGHT - 1) as usize {
226            Err(MerkleTreeError::IncorrectPathLength(path.len()))?
227        } else {
228            Ok(MerkleTreePath {
229                path,
230            })
231        }
232    }
233}
234
235#[derive(Debug)]
236pub enum MerkleTreeError {
237    IncorrectLeafIndex(usize),
238    IncorrectPathLength(usize),
239}
240
241impl std::fmt::Display for MerkleTreeError {
242    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
243        let msg = match self {
244            MerkleTreeError::IncorrectLeafIndex(index) => format!("incorrect leaf index: {}", index),
245            MerkleTreeError::IncorrectPathLength(len) => format!("incorrect path length: {}", len),
246        };
247        write!(f, "{}", msg)
248    }
249}
250
251impl std::error::Error for MerkleTreeError {
252    #[inline]
253    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
254        None
255    }
256}
257
258/// Returns the log2 value of the given number.
259#[inline]
260fn log2(number: usize) -> usize {
261    (number as f64).log2() as usize
262}
263
264/// Returns the height of the tree, given the size of the tree.
265#[inline]
266fn tree_height(tree_size: usize) -> usize {
267    log2(tree_size + 1)
268}
269
270/// Returns true iff the index represents the root.
271#[inline]
272fn is_root(index: usize) -> bool {
273    index == 0
274}
275
276/// Returns the index of the left child, given an index.
277#[inline]
278fn left_child(index: usize) -> usize {
279    2 * index + 1
280}
281
282/// Returns the index of the right child, given an index.
283#[inline]
284fn right_child(index: usize) -> usize {
285    2 * index + 2
286}
287
288/// Returns the index of the sibling, given an index.
289#[inline]
290fn sibling(index: usize) -> Option<usize> {
291    if index == 0 {
292        None
293    } else if is_left_child(index) {
294        Some(index + 1)
295    } else {
296        Some(index - 1)
297    }
298}
299
300/// Returns true iff the given index represents a left child.
301#[inline]
302fn is_left_child(index: usize) -> bool {
303    index % 2 == 1
304}
305
306/// Returns the index of the parent, given an index.
307#[inline]
308fn parent(index: usize) -> Option<usize> {
309    if index > 0 {
310        Some((index - 1) >> 1)
311    } else {
312        None
313    }
314}
315
316#[inline]
317fn convert_index_to_last_level(index: usize, tree_height: usize) -> usize {
318    index + (1 << (tree_height - 1)) - 1
319}
320
321/// Returns the output hash, given a left and right hash value.
322pub(crate) fn hash_inner_node<H: FixedLengthCRH>(
323    parameters: &H::Parameters,
324    left: &H::Output,
325    right: &H::Output,
326    buffer: &mut [u8],
327) -> Result<H::Output, Error> {
328    use std::io::Cursor;
329    let mut writer = Cursor::new(buffer);
330    // Construct left input.
331    left.write(&mut writer)?;
332
333    // Construct right input.
334    right.write(&mut writer)?;
335
336    let buffer = writer.into_inner();
337    H::evaluate(parameters, &buffer[..(H::INPUT_SIZE_BITS / 8)])
338}
339
340/// Returns the hash of a leaf.
341pub(crate) fn hash_leaf<H: FixedLengthCRH, L: ToBytes>(
342    parameters: &H::Parameters,
343    leaf: &L,
344    buffer: &mut [u8],
345) -> Result<H::Output, Error> {
346    use std::io::Cursor;
347    let mut writer = Cursor::new(buffer);
348    leaf.write(&mut writer)?;
349
350    let buffer = writer.into_inner();
351    H::evaluate(parameters, &buffer[..(H::INPUT_SIZE_BITS / 8)])
352}
353
354pub(crate) fn hash_empty<H: FixedLengthCRH>(
355    parameters: &H::Parameters,
356) -> Result<H::Output, Error> {
357    let empty_buffer = vec![0u8; H::INPUT_SIZE_BITS / 8];
358    H::evaluate(parameters, &empty_buffer)
359}
360
361
362#[cfg(test)]
363mod test {
364    use crate::{crh::{pedersen::*, *}, merkle_tree::*};
365    use algebra::curves::jubjub::JubJubAffine as JubJub;
366    use rand::SeedableRng;
367    use rand_xorshift::XorShiftRng;
368
369    #[derive(Clone)]
370    pub(super) struct Window4x256;
371    impl PedersenWindow for Window4x256 {
372        const WINDOW_SIZE: usize = 4;
373        const NUM_WINDOWS: usize = 256;
374    }
375
376
377    type H = PedersenCRH<JubJub, Window4x256>;
378
379    struct JubJubMerkleTreeParams;
380    
381    impl MerkleTreeConfig for JubJubMerkleTreeParams {
382        const HEIGHT: usize = 32;
383        type H = H;
384    }
385    type JubJubMerkleTree = MerkleHashTree<JubJubMerkleTreeParams>;
386
387
388    fn generate_merkle_tree<L: ToBytes + Clone + Eq>(leaves: &[L]) -> () {
389        let mut rng = XorShiftRng::seed_from_u64(9174123u64);
390
391        let crh_parameters = Rc::new(H::setup(&mut rng).unwrap());
392        let tree = JubJubMerkleTree::new(crh_parameters.clone(), &leaves).unwrap();
393        let root = tree.root();
394        for (i, leaf) in leaves.iter().enumerate() {
395            let proof = tree.generate_proof(i, &leaf).unwrap();
396            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
397        }
398    }
399
400    #[test]
401    fn good_root_test() {
402        let mut leaves = Vec::new();
403        for i in 0..4u8 {
404            leaves.push([i, i, i, i, i, i, i, i]);
405        }
406        generate_merkle_tree(&leaves);
407        let mut leaves = Vec::new();
408        for i in 0..100u8 {
409            leaves.push([i, i, i, i, i, i, i, i]);
410        }
411        generate_merkle_tree(&leaves);
412    }
413
414    fn bad_merkle_tree_verify<L: ToBytes + Clone + Eq>(leaves: &[L]) -> () {
415        use algebra::groups::Group;
416        let mut rng = XorShiftRng::seed_from_u64(13423423u64);
417
418        let crh_parameters = Rc::new(H::setup(&mut rng).unwrap());
419        let tree = JubJubMerkleTree::new(crh_parameters.clone(), &leaves).unwrap();
420        let root = JubJub::zero();
421        for (i, leaf) in leaves.iter().enumerate() {
422            let proof = tree.generate_proof(i, &leaf).unwrap();
423            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
424        }
425    }
426
427    #[should_panic]
428    #[test]
429    fn bad_root_test() {
430        let mut leaves = Vec::new();
431        for i in 0..4u8 {
432            leaves.push([i, i, i, i, i, i, i, i]);
433        }
434        generate_merkle_tree(&leaves);
435        let mut leaves = Vec::new();
436        for i in 0..100u8 {
437            leaves.push([i, i, i, i, i, i, i, i]);
438        }
439        bad_merkle_tree_verify(&leaves);
440    }
441}