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