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