1use super::{BTreeMap, MerkleError, Rpo256, RpoDigest, Vec, Word};
  2
  3#[cfg(test)]
  4mod tests;
  5
  6// SPARSE MERKLE TREE
  7// ================================================================================================
  8
  9/// A sparse Merkle tree with 63-bit keys and 4-element leaf values, without compaction.
 10/// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`.
 11/// The root of the tree is recomputed on each new leaf update.
 12#[derive(Debug, Clone, PartialEq, Eq)]
 13pub struct SimpleSmt {
 14    root: Word,
 15    depth: u32,
 16    store: Store,
 17}
 18
 19impl SimpleSmt {
 20    // CONSTANTS
 21    // --------------------------------------------------------------------------------------------
 22
 23    /// Minimum supported depth.
 24    pub const MIN_DEPTH: u32 = 1;
 25
 26    /// Maximum supported depth.
 27    pub const MAX_DEPTH: u32 = 63;
 28
 29    // CONSTRUCTORS
 30    // --------------------------------------------------------------------------------------------
 31
 32    /// Creates a new simple SMT.
 33    ///
 34    /// The provided entries will be tuples of the leaves and their corresponding keys.
 35    ///
 36    /// # Errors
 37    ///
 38    /// The function will fail if the provided entries count exceed the maximum tree capacity, that
 39    /// is `2^{depth}`.
 40    pub fn new<R, I>(entries: R, depth: u32) -> Result<Self, MerkleError>
 41    where
 42        R: IntoIterator<IntoIter = I>,
 43        I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
 44    {
 45        let mut entries = entries.into_iter();
 46
 47        // validate the range of the depth.
 48        let max = 1 << depth;
 49        if depth < Self::MIN_DEPTH {
 50            return Err(MerkleError::DepthTooSmall(depth));
 51        } else if Self::MAX_DEPTH < depth {
 52            return Err(MerkleError::DepthTooBig(depth));
 53        } else if entries.len() > max {
 54            return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
 55        }
 56
 57        let (store, root) = Store::new(depth);
 58        let mut tree = Self { root, depth, store };
 59        entries.try_for_each(|(key, leaf)| tree.insert_leaf(key, leaf))?;
 60
 61        Ok(tree)
 62    }
 63
 64    /// Returns the root of this Merkle tree.
 65    pub const fn root(&self) -> Word {
 66        self.root
 67    }
 68
 69    /// Returns the depth of this Merkle tree.
 70    pub const fn depth(&self) -> u32 {
 71        self.depth
 72    }
 73
 74    /// Returns the set count of the keys of the leaves.
 75    pub fn leaves_count(&self) -> usize {
 76        self.store.leaves_count()
 77    }
 78
 79    /// Returns a node at the specified key
 80    ///
 81    /// # Errors
 82    /// Returns an error if:
 83    /// * The specified depth is greater than the depth of the tree.
 84    /// * The specified key does not exist
 85    pub fn get_node(&self, depth: u32, key: u64) -> Result<Word, MerkleError> {
 86        if depth == 0 {
 87            Err(MerkleError::DepthTooSmall(depth))
 88        } else if depth > self.depth() {
 89            Err(MerkleError::DepthTooBig(depth))
 90        } else if depth == self.depth() {
 91            self.store.get_leaf_node(key)
 92        } else {
 93            let branch_node = self.store.get_branch_node(key, depth)?;
 94            Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
 95        }
 96    }
 97
 98    /// Returns a Merkle path from the node at the specified key to the root. The node itself is
 99    /// not included in the path.
100    ///
101    /// # Errors
102    /// Returns an error if:
103    /// * The specified key does not exist as a branch or leaf node
104    /// * The specified depth is greater than the depth of the tree.
105    pub fn get_path(&self, depth: u32, key: u64) -> Result<Vec<Word>, MerkleError> {
106        if depth == 0 {
107            return Err(MerkleError::DepthTooSmall(depth));
108        } else if depth > self.depth() {
109            return Err(MerkleError::DepthTooBig(depth));
110        } else if depth == self.depth() && !self.store.check_leaf_node_exists(key) {
111            return Err(MerkleError::InvalidIndex(self.depth(), key));
112        }
113
114        let mut path = Vec::with_capacity(depth as usize);
115        let mut curr_key = key;
116        for n in (0..depth).rev() {
117            let parent_key = curr_key >> 1;
118            let parent_node = self.store.get_branch_node(parent_key, n)?;
119            let sibling_node = if curr_key & 1 == 1 {
120                parent_node.left
121            } else {
122                parent_node.right
123            };
124            path.push(sibling_node.into());
125            curr_key >>= 1;
126        }
127        Ok(path)
128    }
129
130    /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
131    /// included in the path.
132    ///
133    /// # Errors
134    /// Returns an error if:
135    /// * The specified key does not exist as a leaf node.
136    pub fn get_leaf_path(&self, key: u64) -> Result<Vec<Word>, MerkleError> {
137        self.get_path(self.depth(), key)
138    }
139
140    /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
141    ///
142    /// # Errors
143    /// Returns an error if the specified key is not a valid leaf index for this tree.
144    pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
145        if !self.store.check_leaf_node_exists(key) {
146            return Err(MerkleError::InvalidIndex(self.depth(), key));
147        }
148        self.insert_leaf(key, value)?;
149
150        Ok(())
151    }
152
153    /// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree
154    pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
155        self.store.insert_leaf_node(key, value);
156
157        let depth = self.depth();
158        let mut curr_key = key;
159        let mut curr_node: RpoDigest = value.into();
160        for n in (0..depth).rev() {
161            let parent_key = curr_key >> 1;
162            let parent_node = self
163                .store
164                .get_branch_node(parent_key, n)
165                .unwrap_or_else(|_| self.store.get_empty_node((n + 1) as usize));
166            let (left, right) = if curr_key & 1 == 1 {
167                (parent_node.left, curr_node)
168            } else {
169                (curr_node, parent_node.right)
170            };
171
172            self.store.insert_branch_node(parent_key, n, left, right);
173            curr_key = parent_key;
174            curr_node = Rpo256::merge(&[left, right]);
175        }
176        self.root = curr_node.into();
177
178        Ok(())
179    }
180}
181
182// STORE
183// ================================================================================================
184
185/// A data store for sparse Merkle tree key-value pairs.
186/// Leaves and branch nodes are stored separately in B-tree maps, indexed by key and (key, depth)
187/// respectively. Hashes for blank subtrees at each layer are stored in `empty_hashes`, beginning
188/// with the root hash of an empty tree, and ending with the zero value of a leaf node.
189#[derive(Debug, Clone, PartialEq, Eq)]
190struct Store {
191    branches: BTreeMap<(u64, u32), BranchNode>,
192    leaves: BTreeMap<u64, Word>,
193    empty_hashes: Vec<RpoDigest>,
194    depth: u32,
195}
196
197#[derive(Debug, Default, Clone, PartialEq, Eq)]
198struct BranchNode {
199    left: RpoDigest,
200    right: RpoDigest,
201}
202
203impl Store {
204    fn new(depth: u32) -> (Self, Word) {
205        let branches = BTreeMap::new();
206        let leaves = BTreeMap::new();
207
208        // Construct empty node digests for each layer of the tree
209        let empty_hashes: Vec<RpoDigest> = (0..depth + 1)
210            .scan(Word::default().into(), |state, _| {
211                let value = *state;
212                *state = Rpo256::merge(&[value, value]);
213                Some(value)
214            })
215            .collect::<Vec<_>>()
216            .into_iter()
217            .rev()
218            .collect();
219
220        let root = empty_hashes[0].into();
221        let store = Self {
222            branches,
223            leaves,
224            empty_hashes,
225            depth,
226        };
227
228        (store, root)
229    }
230
231    fn get_empty_node(&self, depth: usize) -> BranchNode {
232        let digest = self.empty_hashes[depth];
233        BranchNode {
234            left: digest,
235            right: digest,
236        }
237    }
238
239    fn check_leaf_node_exists(&self, key: u64) -> bool {
240        self.leaves.contains_key(&key)
241    }
242
243    fn get_leaf_node(&self, key: u64) -> Result<Word, MerkleError> {
244        self.leaves
245            .get(&key)
246            .cloned()
247            .ok_or(MerkleError::InvalidIndex(self.depth, key))
248    }
249
250    fn insert_leaf_node(&mut self, key: u64, node: Word) {
251        self.leaves.insert(key, node);
252    }
253
254    fn get_branch_node(&self, key: u64, depth: u32) -> Result<BranchNode, MerkleError> {
255        self.branches
256            .get(&(key, depth))
257            .cloned()
258            .ok_or(MerkleError::InvalidIndex(depth, key))
259    }
260
261    fn insert_branch_node(&mut self, key: u64, depth: u32, left: RpoDigest, right: RpoDigest) {
262        let node = BranchNode { left, right };
263        self.branches.insert((key, depth), node);
264    }
265
266    fn leaves_count(&self) -> usize {
267        self.leaves.len()
268    }
269}