update

src/data.rs

@@ -1,16 +1,27 @@
-use core::{cell::RefCell, iter::Chain};
-
 use super::utils::collections::{btree_map, BTreeMap, BTreeSet};
+use core::{cell::RefCell, iter::Chain};
 
+/// A [DataRecorder] that records read requests to the underlying key-value map.
+/// The data recorder is used to generate a proof for read requests. The proof is used to
+/// enable stateless execution of programs on the Miden virtual machine.
+///
+/// The [DataRecorder] is composed of three parts:
+/// - `init`: which contains the initial key-value pairs from the underlying data set.
+/// - `new`: which contains key-value pairs which are created during transaction execution.
+/// - `recorder`: which contains the keys from the initial data set that are read during
+///   program execution.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct DataRecorder<K, V> {
-    // TODO: use smart pointers
     init: BTreeMap<K, V>,
     new: BTreeMap<K, V>,
     recorder: RefCell<BTreeSet<K>>,
 }
 
 impl<K: Ord + Clone, V: Clone> DataRecorder<K, V> {
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+    /// Returns a new [DataRecorder] instance initialized with the provided key-value pairs
+    /// ([BTreeMap]).
     pub fn new(init: BTreeMap<K, V>) -> Self {
         DataRecorder {
             init,
@@ -19,6 +30,10 @@ impl<K: Ord + Clone, V: Clone> DataRecorder<K, V> {
         }
     }
 
+    /// ACCESSORS
+    /// --------------------------------------------------------------------------------------------
+    /// Returns a reference to the value associated with the given key if the value exists. If the
+    /// key is part of the initial data set, the key access is recorded.
     pub fn get(&self, key: &K) -> Option<&V> {
         if let Some(value) = self.new.get(key) {
             return Some(value);
@@ -33,6 +48,8 @@ impl<K: Ord + Clone, V: Clone> DataRecorder<K, V> {
         }
     }
 
+    /// Returns a boolean to indicate whether the given key exists in the data set. If the key is
+    /// part of the initial data set, the key access is recorded.
     pub fn contains_key(&self, key: &K) -> bool {
         if self.new.contains_key(key) == true {
             return true;
@@ -47,24 +64,53 @@ impl<K: Ord + Clone, V: Clone> DataRecorder<K, V> {
         }
     }
 
+    // TODO: Consider if we are happy with the fact that duplicates are counted twice in both
+    // `len(..)` and `iter(..)` methods.
+
+    /// Returns the number of key-value pairs in the data set. The number includes both the initial
+    /// key-value pairs and the new key-value pairs. It should be noted that duplicates in the initial
+    /// and new data sets will be counted twice.
     pub fn len(&self) -> usize {
         self.init.len() + self.new.len()
     }
 
-    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
-        self.new.insert(key.clone(), value.clone())
-    }
-
+    /// Returns an iterator over the key-value pairs in the data set. The iterator includes both
+    /// the initial key-value pairs and the new key-value pairs. It should be noted that duplicates
+    /// in the initial and new data sets will be included twice.
     pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
         self.init.iter().chain(self.new.iter())
     }
 
+    // MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Inserts a key-value pair into the data set. If the key already exists in the data set, the
+    /// value is updated and the old value is returned.
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        if let Some(value) = self.new.insert(key.clone(), value) {
+            return Some(value);
+        }
+
+        match self.init.get(&key) {
+            None => None,
+            Some(value) => {
+                self.recorder.borrow_mut().insert(key);
+                Some(value.clone())
+            }
+        }
+    }
+
+    /// Extends the data set with the key-value pairs from the provided iterator.
     pub fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
         iter.into_iter().for_each(move |(k, v)| {
             self.insert(k, v);
         });
     }
 
+    // FINALIZER
+    // --------------------------------------------------------------------------------------------
+    /// Consumes the [DataRecorder] and returns a [BTreeMap] containing the key-value pairs from
+    /// the initial data set that were read during recording.
     pub fn into_proof(self) -> BTreeMap<K, V> {
         self.init
             .into_iter()
@@ -73,12 +119,14 @@ impl<K: Ord + Clone, V: Clone> DataRecorder<K, V> {
     }
 }
 
+/// Default implementation for [DataRecorder].
 impl<K: Ord + Clone, V: Clone> Default for DataRecorder<K, V> {
     fn default() -> Self {
         DataRecorder::new(BTreeMap::new())
     }
 }
 
+/// Implementation of [IntoIterator] for [DataRecorder].
 impl<K, V> IntoIterator for DataRecorder<K, V> {
     type Item = (K, V);
     type IntoIter = Chain<btree_map::IntoIter<K, V>, btree_map::IntoIter<K, V>>;
@@ -87,8 +135,3 @@ impl<K, V> IntoIterator for DataRecorder<K, V> {
         self.init.into_iter().chain(self.new.into_iter())
     }
 }
-
-#[derive(Debug)]
-pub enum DataBackendError {
-    RecorderNotSet,
-}

src/merkle/mod.rs

@@ -58,7 +58,6 @@ pub enum MerkleError {
     NodeNotInStore(Word, NodeIndex),
     NumLeavesNotPowerOfTwo(usize),
     RootNotInStore(Word),
-    RecorderNotSet,
 }
 
 impl fmt::Display for MerkleError {
@@ -86,7 +85,6 @@ impl fmt::Display for MerkleError {
                 write!(f, "the leaves count {leaves} is not a power of 2")
             }
             RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root),
-            RecorderNotSet => write!(f, "the recorder is not set"),
         }
     }
 }

src/merkle/store/mod.rs

@@ -267,7 +267,7 @@ impl MerkleStore {
     /// nodes which are descendants of the specified roots.
     ///
     /// The roots for which no descendants exist in this Merkle store are ignored.
-    pub fn subset<I, R>(&mut self, roots: I) -> MerkleStore
+    pub fn subset<I, R>(&self, roots: I) -> MerkleStore
     where
         I: Iterator<Item = R>,
         R: Borrow<Word>,

src/merkle/store/recorder.rs

@@ -5,17 +5,18 @@ use super::{
 };
 use core::borrow::Borrow;
 
-/// An in-memory data store for Merkelized data.
+/// An in-memory data store for Merkelized data with data access recording capabilities.
 ///
 /// This is a in memory data store for Merkle trees, this store allows all the nodes of multiple
 /// trees to live as long as necessary and without duplication, this allows the implementation of
-/// space efficient persistent data structures.
+/// space efficient persistent data structures. The store also allows recording of data access
+/// such that a proof of data access can be generated.
 ///
 /// Example usage:
 ///
 /// ```rust
 /// # use miden_crypto::{ZERO, Felt, Word};
-/// # use miden_crypto::merkle::{NodeIndex, MerkleStore, MerkleTree};
+/// # use miden_crypto::merkle::{NodeIndex, RecordingMerkleStore, MerkleTree};
 /// # use miden_crypto::hash::rpo::Rpo256;
 /// # const fn int_to_node(value: u64) -> Word {
 /// #     [Felt::new(value), ZERO, ZERO, ZERO]
@@ -33,7 +34,7 @@ use core::borrow::Borrow;
 /// # let T1 = MerkleTree::new([A, B, C, D, E, F, G, H1].to_vec()).expect("even number of leaves provided");
 /// # let ROOT0 = T0.root();
 /// # let ROOT1 = T1.root();
-/// let mut store = MerkleStore::new();
+/// let mut store = RecordingMerkleStore::new();
 ///
 /// // the store is initialized with the SMT empty nodes
 /// assert_eq!(store.num_internal_nodes(), 255);
@@ -67,6 +68,9 @@ use core::borrow::Borrow;
 /// // Common internal nodes are shared, the two added trees have a total of 30, but the store has
 /// // only 10 new entries, corresponding to the 10 unique internal nodes of these trees.
 /// assert_eq!(store.num_internal_nodes() - 255, 10);
+///
+/// // Convert the store into a proof
+/// let proof = store.into_proof();
 /// ```
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct RecordingMerkleStore {
@@ -83,7 +87,7 @@ impl RecordingMerkleStore {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
-    /// Creates an empty `MerkleStore` instance.
+    /// Creates an empty [RecordingMerkleStore] instance.
     pub fn new() -> RecordingMerkleStore {
         // pre-populate the store with the empty hashes
         let subtrees = EmptySubtreeRoots::empty_hashes(255);
@@ -462,6 +466,3 @@ impl Extend<InnerNodeInfo> for RecordingMerkleStore {
         self.extend(iter.into_iter());
     }
 }
-
-// SERIALIZATION
-// ================================================================================================

src/merkle/store/tests.rs

@@ -752,31 +752,31 @@ fn mstore_subset() {
 
     // --- extract all 3 subtrees ---------------------------------------------
 
-    let mut substore = store.subset([subtree1.root(), subtree2.root(), subtree3.root()].iter());
+    let substore = store.subset([subtree1.root(), subtree2.root(), subtree3.root()].iter());
 
     // number of nodes should increase by 4: 3 nodes form subtree1 and 1 node from subtree3
     assert_eq!(substore.nodes.len(), empty_store_num_nodes + 4);
 
     // make sure paths that all subtrees are in the store
-    check_mstore_subtree(&mut substore, &subtree1);
-    check_mstore_subtree(&mut substore, &subtree2);
-    check_mstore_subtree(&mut substore, &subtree3);
+    check_mstore_subtree(&substore, &subtree1);
+    check_mstore_subtree(&substore, &subtree2);
+    check_mstore_subtree(&substore, &subtree3);
 
     // --- extract subtrees 1 and 3 -------------------------------------------
     // this should give the same result as above as subtree2 is nested withing subtree1
 
-    let mut substore = store.subset([subtree1.root(), subtree3.root()].iter());
+    let substore = store.subset([subtree1.root(), subtree3.root()].iter());
 
     // number of nodes should increase by 4: 3 nodes form subtree1 and 1 node from subtree3
     assert_eq!(substore.nodes.len(), empty_store_num_nodes + 4);
 
     // make sure paths that all subtrees are in the store
-    check_mstore_subtree(&mut substore, &subtree1);
-    check_mstore_subtree(&mut substore, &subtree2);
-    check_mstore_subtree(&mut substore, &subtree3);
+    check_mstore_subtree(&substore, &subtree1);
+    check_mstore_subtree(&substore, &subtree2);
+    check_mstore_subtree(&substore, &subtree3);
 }
 
-fn check_mstore_subtree(store: &mut MerkleStore, subtree: &MerkleTree) {
+fn check_mstore_subtree(store: &MerkleStore, subtree: &MerkleTree) {
     for (i, value) in subtree.leaves() {
         let index = NodeIndex::new(subtree.depth(), i).unwrap();
         let path1 = store.get_path(subtree.root(), index).unwrap();