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