1use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, 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: u8,
 16    store: Store,
 17}
 18
 19impl SimpleSmt {
 20    // CONSTANTS
 21    // --------------------------------------------------------------------------------------------
 22
 23    /// Minimum supported depth.
 24    pub const MIN_DEPTH: u8 = 1;
 25
 26    /// Maximum supported depth.
 27    pub const MAX_DEPTH: u8 = 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: u8) -> 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 as u64));
 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) -> u8 {
 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, index: &NodeIndex) -> Result<Word, MerkleError> {
 86        if index.is_root() {
 87            Err(MerkleError::DepthTooSmall(index.depth()))
 88        } else if index.depth() > self.depth() {
 89            Err(MerkleError::DepthTooBig(index.depth() as u64))
 90        } else if index.depth() == self.depth() {
 91            self.store.get_leaf_node(index.value())
 92        } else {
 93            let branch_node = self.store.get_branch_node(index)?;
 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, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
106        if index.is_root() {
107            return Err(MerkleError::DepthTooSmall(index.depth()));
108        } else if index.depth() > self.depth() {
109            return Err(MerkleError::DepthTooBig(index.depth() as u64));
110        } else if index.depth() == self.depth() && !self.store.check_leaf_node_exists(index.value())
111        {
112            return Err(MerkleError::InvalidIndex(index.with_depth(self.depth())));
113        }
114
115        let mut path = Vec::with_capacity(index.depth() as usize);
116        for _ in 0..index.depth() {
117            let is_right = index.is_value_odd();
118            index.move_up();
119            let BranchNode { left, right } = self.store.get_branch_node(&index)?;
120            let value = if is_right { left } else { right };
121            path.push(*value);
122        }
123        Ok(path.into())
124    }
125
126    /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
127    /// included in the path.
128    ///
129    /// # Errors
130    /// Returns an error if:
131    /// * The specified key does not exist as a leaf node.
132    pub fn get_leaf_path(&self, key: u64) -> Result<MerklePath, MerkleError> {
133        self.get_path(NodeIndex::new(self.depth(), key))
134    }
135
136    /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
137    ///
138    /// # Errors
139    /// Returns an error if the specified key is not a valid leaf index for this tree.
140    pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
141        if !self.store.check_leaf_node_exists(key) {
142            return Err(MerkleError::InvalidIndex(NodeIndex::new(self.depth(), key)));
143        }
144        self.insert_leaf(key, value)?;
145
146        Ok(())
147    }
148
149    /// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree
150    pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
151        self.store.insert_leaf_node(key, value);
152
153        // TODO consider using a map `index |-> word` instead of `index |-> (word, word)`
154        let mut index = NodeIndex::new(self.depth(), key);
155        let mut value = RpoDigest::from(value);
156        for _ in 0..index.depth() {
157            let is_right = index.is_value_odd();
158            index.move_up();
159            let BranchNode { left, right } = self
160                .store
161                .get_branch_node(&index)
162                .unwrap_or_else(|_| self.store.get_empty_node(index.depth() as usize + 1));
163            let (left, right) = if is_right {
164                (left, value)
165            } else {
166                (value, right)
167            };
168            self.store.insert_branch_node(index, left, right);
169            value = Rpo256::merge(&[left, right]);
170        }
171        self.root = value.into();
172        Ok(())
173    }
174}
175
176// STORE
177// ================================================================================================
178
179/// A data store for sparse Merkle tree key-value pairs.
180/// Leaves and branch nodes are stored separately in B-tree maps, indexed by key and (key, depth)
181/// respectively. Hashes for blank subtrees at each layer are stored in `empty_hashes`, beginning
182/// with the root hash of an empty tree, and ending with the zero value of a leaf node.
183#[derive(Debug, Clone, PartialEq, Eq)]
184struct Store {
185    branches: BTreeMap<NodeIndex, BranchNode>,
186    leaves: BTreeMap<u64, Word>,
187    empty_hashes: Vec<RpoDigest>,
188    depth: u8,
189}
190
191#[derive(Debug, Default, Clone, PartialEq, Eq)]
192struct BranchNode {
193    left: RpoDigest,
194    right: RpoDigest,
195}
196
197impl Store {
198    fn new(depth: u8) -> (Self, Word) {
199        let branches = BTreeMap::new();
200        let leaves = BTreeMap::new();
201
202        // Construct empty node digests for each layer of the tree
203        let empty_hashes: Vec<RpoDigest> = (0..depth + 1)
204            .scan(Word::default().into(), |state, _| {
205                let value = *state;
206                *state = Rpo256::merge(&[value, value]);
207                Some(value)
208            })
209            .collect::<Vec<_>>()
210            .into_iter()
211            .rev()
212            .collect();
213
214        let root = empty_hashes[0].into();
215        let store = Self {
216            branches,
217            leaves,
218            empty_hashes,
219            depth,
220        };
221
222        (store, root)
223    }
224
225    fn get_empty_node(&self, depth: usize) -> BranchNode {
226        let digest = self.empty_hashes[depth];
227        BranchNode {
228            left: digest,
229            right: digest,
230        }
231    }
232
233    fn check_leaf_node_exists(&self, key: u64) -> bool {
234        self.leaves.contains_key(&key)
235    }
236
237    fn get_leaf_node(&self, key: u64) -> Result<Word, MerkleError> {
238        self.leaves
239            .get(&key)
240            .cloned()
241            .ok_or(MerkleError::InvalidIndex(NodeIndex::new(self.depth, key)))
242    }
243
244    fn insert_leaf_node(&mut self, key: u64, node: Word) {
245        self.leaves.insert(key, node);
246    }
247
248    fn get_branch_node(&self, index: &NodeIndex) -> Result<BranchNode, MerkleError> {
249        self.branches
250            .get(index)
251            .cloned()
252            .ok_or(MerkleError::InvalidIndex(*index))
253    }
254
255    fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {
256        let branch = BranchNode { left, right };
257        self.branches.insert(index, branch);
258    }
259
260    fn leaves_count(&self) -> usize {
261        self.leaves.len()
262    }
263}