arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/merkle/smt/mod.rs
src/merkle/smt/mod.rsRaw
  1use alloc::vec::Vec;
  2
  3use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
  4use crate::{
  5    hash::rpo::{Rpo256, RpoDigest},
  6    Felt, Word, EMPTY_WORD,
  7};
  8
  9mod full;
 10pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};
 11
 12mod simple;
 13pub use simple::SimpleSmt;
 14
 15// CONSTANTS
 16// ================================================================================================
 17
 18/// Minimum supported depth.
 19pub const SMT_MIN_DEPTH: u8 = 1;
 20
 21/// Maximum supported depth.
 22pub const SMT_MAX_DEPTH: u8 = 64;
 23
 24// SPARSE MERKLE TREE
 25// ================================================================================================
 26
 27/// An abstract description of a sparse Merkle tree.
 28///
 29/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
 30/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
 31/// value was not explicitly set, then its value is the default value. Typically, the vast majority
 32/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
 33/// representation of the tree will only keep track of the leaves that have a different value from
 34/// the default.
 35///
 36/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
 37/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
 38/// tree is just a single value, and is probably a programming mistake.
 39///
 40/// Every key maps to one leaf. If there are as many keys as there are leaves, then
 41/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
 42/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
 43/// must accomodate all keys that map to the same leaf.
 44///
 45/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
 46pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
 47    /// The type for a key
 48    type Key: Clone;
 49    /// The type for a value
 50    type Value: Clone + PartialEq;
 51    /// The type for a leaf
 52    type Leaf;
 53    /// The type for an opening (i.e. a "proof") of a leaf
 54    type Opening;
 55
 56    /// The default value used to compute the hash of empty leaves
 57    const EMPTY_VALUE: Self::Value;
 58
 59    // PROVIDED METHODS
 60    // ---------------------------------------------------------------------------------------------
 61
 62    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
 63    /// path to the leaf, as well as the leaf itself.
 64    fn open(&self, key: &Self::Key) -> Self::Opening {
 65        let leaf = self.get_leaf(key);
 66
 67        let mut index: NodeIndex = {
 68            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
 69            leaf_index.into()
 70        };
 71
 72        let merkle_path = {
 73            let mut path = Vec::with_capacity(index.depth() as usize);
 74            for _ in 0..index.depth() {
 75                let is_right = index.is_value_odd();
 76                index.move_up();
 77                let InnerNode { left, right } = self.get_inner_node(index);
 78                let value = if is_right { left } else { right };
 79                path.push(value);
 80            }
 81
 82            MerklePath::new(path)
 83        };
 84
 85        Self::path_and_leaf_to_opening(merkle_path, leaf)
 86    }
 87
 88    /// Inserts a value at the specified key, returning the previous value associated with that key.
 89    /// Recall that by definition, any key that hasn't been updated is associated with
 90    /// [`Self::EMPTY_VALUE`].
 91    ///
 92    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
 93    /// updating the root itself.
 94    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
 95        let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);
 96
 97        // if the old value and new value are the same, there is nothing to update
 98        if value == old_value {
 99            return value;
100        }
101
102        let leaf = self.get_leaf(&key);
103        let node_index = {
104            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
105            leaf_index.into()
106        };
107
108        self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
109
110        old_value
111    }
112
113    /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
114    /// `node_hash_at_index` is the hash of the node stored at index.
115    fn recompute_nodes_from_index_to_root(
116        &mut self,
117        mut index: NodeIndex,
118        node_hash_at_index: RpoDigest,
119    ) {
120        let mut node_hash = node_hash_at_index;
121        for node_depth in (0..index.depth()).rev() {
122            let is_right = index.is_value_odd();
123            index.move_up();
124            let InnerNode { left, right } = self.get_inner_node(index);
125            let (left, right) = if is_right {
126                (left, node_hash)
127            } else {
128                (node_hash, right)
129            };
130            node_hash = Rpo256::merge(&[left, right]);
131
132            if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
133                // If a subtree is empty, when can remove the inner node, since it's equal to the
134                // default value
135                self.remove_inner_node(index)
136            } else {
137                self.insert_inner_node(index, InnerNode { left, right });
138            }
139        }
140        self.set_root(node_hash);
141    }
142
143    // REQUIRED METHODS
144    // ---------------------------------------------------------------------------------------------
145
146    /// The root of the tree
147    fn root(&self) -> RpoDigest;
148
149    /// Sets the root of the tree
150    fn set_root(&mut self, root: RpoDigest);
151
152    /// Retrieves an inner node at the given index
153    fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
154
155    /// Inserts an inner node at the given index
156    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode);
157
158    /// Removes an inner node at the given index
159    fn remove_inner_node(&mut self, index: NodeIndex);
160
161    /// Inserts a leaf node, and returns the value at the key if already exists
162    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
163
164    /// Returns the leaf at the specified index.
165    fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
166
167    /// Returns the hash of a leaf
168    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
169
170    /// Maps a key to a leaf index
171    fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
172
173    /// Maps a (MerklePath, Self::Leaf) to an opening.
174    ///
175    /// The length `path` is guaranteed to be equal to `DEPTH`
176    fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
177}
178
179// INNER NODE
180// ================================================================================================
181
182#[derive(Debug, Default, Clone, PartialEq, Eq)]
183#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
184pub(crate) struct InnerNode {
185    pub left: RpoDigest,
186    pub right: RpoDigest,
187}
188
189impl InnerNode {
190    pub fn hash(&self) -> RpoDigest {
191        Rpo256::merge(&[self.left, self.right])
192    }
193}
194
195// LEAF INDEX
196// ================================================================================================
197
198/// The index of a leaf, at a depth known at compile-time.
199#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
200#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
201pub struct LeafIndex<const DEPTH: u8> {
202    index: NodeIndex,
203}
204
205impl<const DEPTH: u8> LeafIndex<DEPTH> {
206    pub fn new(value: u64) -> Result<Self, MerkleError> {
207        if DEPTH < SMT_MIN_DEPTH {
208            return Err(MerkleError::DepthTooSmall(DEPTH));
209        }
210
211        Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
212    }
213
214    pub fn value(&self) -> u64 {
215        self.index.value()
216    }
217}
218
219impl LeafIndex<SMT_MAX_DEPTH> {
220    pub const fn new_max_depth(value: u64) -> Self {
221        LeafIndex {
222            index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
223        }
224    }
225}
226
227impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
228    fn from(value: LeafIndex<DEPTH>) -> Self {
229        value.index
230    }
231}
232
233impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
234    type Error = MerkleError;
235
236    fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
237        if node_index.depth() != DEPTH {
238            return Err(MerkleError::InvalidDepth {
239                expected: DEPTH,
240                provided: node_index.depth(),
241            });
242        }
243
244        Self::new(node_index.value())
245    }
246}