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