light_merkle_tree_reference/
indexed.rs

1use std::{fmt::Debug, marker::PhantomData};
2
3use light_hasher::{bigint::bigint_to_be_bytes_array, Hasher, HasherError};
4use light_indexed_array::{
5    array::{IndexedArray, IndexedElement},
6    errors::IndexedArrayError,
7    HIGHEST_ADDRESS_PLUS_ONE,
8};
9use num_bigint::BigUint;
10use num_traits::{CheckedAdd, CheckedSub, Num, ToBytes, Unsigned, Zero};
11use thiserror::Error;
12
13use crate::{MerkleTree, ReferenceMerkleTreeError};
14
15#[derive(Debug, Error, PartialEq)]
16pub enum IndexedReferenceMerkleTreeError {
17    #[error("NonInclusionProofFailedLowerBoundViolated")]
18    NonInclusionProofFailedLowerBoundViolated,
19    #[error("NonInclusionProofFailedHigherBoundViolated")]
20    NonInclusionProofFailedHigherBoundViolated,
21    #[error(transparent)]
22    Indexed(#[from] IndexedArrayError),
23    #[error(transparent)]
24    Reference(#[from] ReferenceMerkleTreeError),
25    #[error(transparent)]
26    Hasher(#[from] HasherError),
27}
28
29#[derive(Debug, Clone)]
30#[repr(C)]
31pub struct IndexedMerkleTree<H, I>
32where
33    H: Hasher,
34    I: CheckedAdd
35        + CheckedSub
36        + Copy
37        + Clone
38        + PartialOrd
39        + ToBytes
40        + TryFrom<usize>
41        + Unsigned
42        + Debug
43        + Into<usize>,
44{
45    pub merkle_tree: MerkleTree<H>,
46    pub indexed_array: IndexedArray<H, I>,
47    _index: PhantomData<I>,
48}
49
50impl<H, I> IndexedMerkleTree<H, I>
51where
52    H: Hasher,
53    I: CheckedAdd
54        + CheckedSub
55        + Copy
56        + Clone
57        + PartialOrd
58        + ToBytes
59        + TryFrom<usize>
60        + Unsigned
61        + Debug
62        + Into<usize>,
63{
64    pub fn new(
65        height: usize,
66        canopy_depth: usize,
67    ) -> Result<Self, IndexedReferenceMerkleTreeError> {
68        let mut merkle_tree = MerkleTree::new(height, canopy_depth);
69
70        let init_next_value = BigUint::from_str_radix(HIGHEST_ADDRESS_PLUS_ONE, 10).unwrap();
71        let indexed_array = IndexedArray::<H, I>::new(BigUint::zero(), init_next_value.clone());
72        let new_leaf = indexed_array
73            .get(0)
74            .unwrap()
75            .hash::<H>(&init_next_value)
76            .unwrap();
77        merkle_tree.append(&new_leaf)?;
78        assert_eq!(merkle_tree.leaf(0), new_leaf);
79        Ok(Self {
80            merkle_tree,
81            indexed_array,
82            _index: PhantomData,
83        })
84    }
85
86    pub fn get_path_of_leaf(
87        &self,
88        index: usize,
89        full: bool,
90    ) -> Result<Vec<[u8; 32]>, ReferenceMerkleTreeError> {
91        self.merkle_tree.get_path_of_leaf(index, full)
92    }
93
94    pub fn get_proof_of_leaf(
95        &self,
96        index: usize,
97        full: bool,
98    ) -> Result<Vec<[u8; 32]>, ReferenceMerkleTreeError> {
99        self.merkle_tree.get_proof_of_leaf(index, full)
100    }
101
102    pub fn root(&self) -> [u8; 32] {
103        self.merkle_tree.root()
104    }
105
106    // TODO: rename input values
107    pub fn update(
108        &mut self,
109        new_low_element: &IndexedElement<I>,
110        new_element: &IndexedElement<I>,
111        new_element_next_value: &BigUint,
112    ) -> Result<(), IndexedReferenceMerkleTreeError> {
113        // Update the low element.
114        let new_low_leaf = new_low_element.hash::<H>(&new_element.value)?;
115        println!("reference update new low leaf hash {:?}", new_low_leaf);
116        self.merkle_tree.update(
117            &new_low_leaf,
118            usize::try_from(new_low_element.index).unwrap(),
119        )?;
120        println!("reference updated root {:?}", self.merkle_tree.root());
121        // Append the new element.
122        let new_leaf = new_element.hash::<H>(new_element_next_value)?;
123        println!("reference update new leaf hash {:?}", new_leaf);
124        self.merkle_tree.append(&new_leaf)?;
125        println!("reference appended root {:?}", self.merkle_tree.root());
126
127        Ok(())
128    }
129
130    // TODO: add append with new value, so that we don't need to compute the lowlevel values manually
131    pub fn append(&mut self, value: &BigUint) -> Result<(), IndexedReferenceMerkleTreeError> {
132        println!("\n\nappending {:?}", value);
133        let nullifier_bundle = self.indexed_array.append(value)?;
134        println!("\n\nnullifier_bundle {:?}", nullifier_bundle);
135        self.update(
136            &nullifier_bundle.new_low_element,
137            &nullifier_bundle.new_element,
138            &nullifier_bundle.new_element_next_value,
139        )?;
140
141        Ok(())
142    }
143
144    pub fn get_non_inclusion_proof(
145        &self,
146        value: &BigUint,
147    ) -> Result<NonInclusionProof, IndexedReferenceMerkleTreeError> {
148        let (low_element, _next_value) =
149            self.indexed_array.find_low_element_for_nonexistent(value)?;
150        let merkle_proof =
151            self.get_proof_of_leaf(usize::try_from(low_element.index).unwrap(), true)?;
152        let higher_range_value = if low_element.next_index() == 0 {
153            self.indexed_array.highest_value.clone()
154        } else {
155            self.indexed_array
156                .get(low_element.next_index())
157                .unwrap()
158                .value
159                .clone()
160        };
161        let non_inclusion_proof = NonInclusionProof {
162            root: self.root(),
163            value: bigint_to_be_bytes_array::<32>(value).unwrap(),
164            leaf_lower_range_value: bigint_to_be_bytes_array::<32>(&low_element.value).unwrap(),
165            leaf_higher_range_value: bigint_to_be_bytes_array::<32>(&higher_range_value).unwrap(),
166            leaf_index: low_element.index.into(),
167            next_index: low_element.next_index(),
168            merkle_proof,
169        };
170        Ok(non_inclusion_proof)
171    }
172
173    pub fn verify_non_inclusion_proof(
174        &self,
175        proof: &NonInclusionProof,
176    ) -> Result<(), IndexedReferenceMerkleTreeError> {
177        let value_big_int = BigUint::from_bytes_be(&proof.value);
178        let lower_end_value = BigUint::from_bytes_be(&proof.leaf_lower_range_value);
179        if lower_end_value >= value_big_int {
180            return Err(IndexedReferenceMerkleTreeError::NonInclusionProofFailedLowerBoundViolated);
181        }
182        let higher_end_value = BigUint::from_bytes_be(&proof.leaf_higher_range_value);
183        if higher_end_value <= value_big_int {
184            return Err(
185                IndexedReferenceMerkleTreeError::NonInclusionProofFailedHigherBoundViolated,
186            );
187        }
188
189        let array_element = IndexedElement::<usize> {
190            value: lower_end_value,
191            index: proof.leaf_index,
192            next_index: proof.next_index,
193        };
194        let leaf_hash = array_element.hash::<H>(&higher_end_value)?;
195        self.merkle_tree
196            .verify(&leaf_hash, &proof.merkle_proof, proof.leaf_index)?;
197        Ok(())
198    }
199}
200
201// TODO: check why next_index is usize while index is I
202/// We prove non-inclusion by:
203/// 1. Showing that value is greater than leaf_lower_range_value and less than leaf_higher_range_value
204/// 2. Showing that the leaf_hash H(leaf_lower_range_value, leaf_next_index, leaf_higher_value) is included in the root (Merkle tree)
205#[derive(Debug)]
206pub struct NonInclusionProof {
207    pub root: [u8; 32],
208    pub value: [u8; 32],
209    pub leaf_lower_range_value: [u8; 32],
210    pub leaf_higher_range_value: [u8; 32],
211    pub leaf_index: usize,
212    pub next_index: usize,
213    pub merkle_proof: Vec<[u8; 32]>,
214}