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