1use winter_math::StarkField;
  2
  3use crate::{
  4    hash::rpo::{Rpo256, RpoDigest},
  5    Word,
  6};
  7
  8use super::{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 value = node_hash_at_index;
122        for _ in 0..index.depth() {
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 { (left, value) } else { (value, right) };
127            self.insert_inner_node(index, InnerNode { left, right });
128            value = Rpo256::merge(&[left, right]);
129        }
130        self.set_root(value);
131    }
132
133    // REQUIRED METHODS
134    // ---------------------------------------------------------------------------------------------
135
136    /// The root of the tree
137    fn root(&self) -> RpoDigest;
138
139    /// Sets the root of the tree
140    fn set_root(&mut self, root: RpoDigest);
141
142    /// Retrieves an inner node at the given index
143    fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
144
145    /// Inserts an inner node at the given index
146    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode);
147
148    /// Inserts a leaf node, and returns the value at the key if already exists
149    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
150
151    /// Returns the leaf at the specified index.
152    fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
153
154    /// Returns the hash of a leaf
155    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
156
157    /// Maps a key to a leaf index
158    fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
159}
160
161// INNER NODE
162// ================================================================================================
163
164#[derive(Debug, Default, Clone, PartialEq, Eq)]
165#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
166pub(crate) struct InnerNode {
167    pub left: RpoDigest,
168    pub right: RpoDigest,
169}
170
171impl InnerNode {
172    pub fn hash(&self) -> RpoDigest {
173        Rpo256::merge(&[self.left, self.right])
174    }
175}
176
177// LEAF INDEX
178// ================================================================================================
179
180/// The index of a leaf, at a depth known at compile-time.
181#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
182#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
183pub struct LeafIndex<const DEPTH: u8> {
184    index: NodeIndex,
185}
186
187impl<const DEPTH: u8> LeafIndex<DEPTH> {
188    pub fn new(value: u64) -> Result<Self, MerkleError> {
189        if DEPTH < SMT_MIN_DEPTH {
190            return Err(MerkleError::DepthTooSmall(DEPTH));
191        }
192
193        Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
194    }
195
196    pub fn value(&self) -> u64 {
197        self.index.value()
198    }
199}
200
201impl LeafIndex<SMT_MAX_DEPTH> {
202    pub const fn new_max_depth(value: u64) -> Self {
203        LeafIndex {
204            index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
205        }
206    }
207}
208
209impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
210    fn from(value: LeafIndex<DEPTH>) -> Self {
211        value.index
212    }
213}
214
215impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
216    type Error = MerkleError;
217
218    fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
219        if node_index.depth() != DEPTH {
220            return Err(MerkleError::InvalidDepth {
221                expected: DEPTH,
222                provided: node_index.depth(),
223            });
224        }
225
226        Self::new(node_index.value())
227    }
228}
229
230impl From<Word> for LeafIndex<SMT_MAX_DEPTH> {
231    fn from(value: Word) -> Self {
232        // We use the most significant `Felt` of a `Word` as the leaf index.
233        Self::new_max_depth(value[3].as_int())
234    }
235}
236
237impl From<RpoDigest> for LeafIndex<SMT_MAX_DEPTH> {
238    fn from(value: RpoDigest) -> Self {
239        Word::from(value).into()
240    }
241}