From 1cdd3dbbfae4c5aa4b8d13e15ccc5d5c7e8b71c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philippe=20Laferri=C3=A8re?= <plafer@protonmail.com>
Date: Mon, 29 Jan 2024 10:13:52 -0500
Subject: [PATCH] Add methods to `Smt` necessary for VM tests (#264)

* Smt::inner_nodes

* Add conversion Smt -> MerkleStore

* add docstring to `Smt`

* add to docstring

* fmt

* add `leaves()` method to `Smt`

* add `kv_pairs` functions

* rewrite `into_elements()` in terms of `into_kv_pairs()`

* change docstring
---
 src/merkle/smt/full/mod.rs | 48 +++++++++++++++++++++++++++++++++++---
 src/merkle/store/mod.rs    |  9 ++++++-
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs
index 0a796a7..cf3b421 100644
--- a/src/merkle/smt/full/mod.rs
+++ b/src/merkle/smt/full/mod.rs
@@ -3,7 +3,7 @@ use core::cmp::Ordering;
 use winter_math::StarkField;
 
 use crate::hash::rpo::Rpo256;
-use crate::merkle::EmptySubtreeRoots;
+use crate::merkle::{EmptySubtreeRoots, InnerNodeInfo};
 use crate::utils::{
     collections::{BTreeMap, BTreeSet, Vec},
     vec,
@@ -25,6 +25,15 @@ pub const SMT_DEPTH: u8 = 64;
 // SMT
 // ================================================================================================
 
+/// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
+/// by 4 field elements.
+///
+/// All leaves sit at depth 64. The most significant element of the key is used to identify the leaf to
+/// which the key maps.
+///
+/// A leaf is either empty, or holds one or more key-value pairs. An empty leaf hashes to the empty
+/// word. Otherwise, a leaf hashes to the hash of its key-value pairs, ordered by key first, value
+/// second.
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct Smt {
@@ -111,6 +120,25 @@ impl Smt {
         <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
     }
 
+    // ITERATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns an iterator over the leaves of this [Smt].
+    pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
+        self.leaves
+            .iter()
+            .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
+    }
+
+    /// Returns an iterator over the inner nodes of this [Smt].
+    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
+        self.inner_nodes.values().map(|e| InnerNodeInfo {
+            value: e.hash(),
+            left: e.left,
+            right: e.right,
+        })
+    }
+
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
@@ -242,10 +270,24 @@ impl SmtLeaf {
 
     /// Converts a leaf to a list of field elements
     pub fn into_elements(self) -> Vec<Felt> {
+        self.into_kv_pairs().into_iter().flat_map(kv_to_elements).collect()
+    }
+
+    /// Returns the key-value pairs in the leaf
+    pub fn kv_pairs(&self) -> Vec<&(RpoDigest, Word)> {
+        match self {
+            SmtLeaf::Empty => Vec::new(),
+            SmtLeaf::Single(kv_pair) => vec![kv_pair],
+            SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(),
+        }
+    }
+
+    /// Converts a leaf the key-value pairs in the leaf
+    pub fn into_kv_pairs(self) -> Vec<(RpoDigest, Word)> {
         match self {
             SmtLeaf::Empty => Vec::new(),
-            SmtLeaf::Single(kv_pair) => kv_to_elements(kv_pair).collect(),
-            SmtLeaf::Multiple(kv_pairs) => kv_pairs.into_iter().flat_map(kv_to_elements).collect(),
+            SmtLeaf::Single(kv_pair) => vec![kv_pair],
+            SmtLeaf::Multiple(kv_pairs) => kv_pairs,
         }
     }
 
diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs
index 5f3e36c..a06e292 100644
--- a/src/merkle/store/mod.rs
+++ b/src/merkle/store/mod.rs
@@ -1,7 +1,7 @@
 use super::{
     mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, KvMap, MerkleError, MerklePath,
     MerkleTree, NodeIndex, PartialMerkleTree, RecordingMap, RootPath, Rpo256, RpoDigest, SimpleSmt,
-    TieredSmt, ValuePath, Vec,
+    Smt, TieredSmt, ValuePath, Vec,
 };
 use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 use core::borrow::Borrow;
@@ -497,6 +497,13 @@ impl<T: KvMap<RpoDigest, StoreNode>, const DEPTH: u8> From<&SimpleSmt<DEPTH>> fo
     }
 }
 
+impl<T: KvMap<RpoDigest, StoreNode>> From<&Smt> for MerkleStore<T> {
+    fn from(value: &Smt) -> Self {
+        let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
+        Self { nodes }
+    }
+}
+
 impl<T: KvMap<RpoDigest, StoreNode>> From<&Mmr> for MerkleStore<T> {
     fn from(value: &Mmr) -> Self {
         let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();