add sorted_pairs_to_leaves() and test for it

src/merkle/smt/full/mod.rs

@@ -351,6 +351,30 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
     fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
         SmtProof::new_unchecked(path, leaf)
     }
+
+    fn pairs_to_leaf(mut pairs: Vec<(RpoDigest, Word)>) -> SmtLeaf {
+        assert!(!pairs.is_empty());
+
+        // FIXME
+        if cfg!(debug_assertions) {
+            let mut control = pairs.clone();
+            control.sort_by_key(|(key, _)| Self::key_to_leaf_index(key).index.value());
+            assert_eq!(control, pairs);
+        }
+
+        if pairs.len() > 1 {
+            SmtLeaf::new_multiple(pairs).unwrap()
+        } else {
+            let (key, value) = pairs.pop().unwrap();
+            // FIXME: should we ever be constructing empty leaves from pairs?
+            if value == Self::EMPTY_VALUE {
+                let index = Self::key_to_leaf_index(&key);
+                SmtLeaf::new_empty(index)
+            } else {
+                SmtLeaf::new_single(key, value)
+            }
+        }
+    }
 }
 
 impl Default for Smt {

src/merkle/smt/mod.rs

@@ -341,11 +341,54 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// Maps a key to a leaf index
     fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
 
+    /// Constructs a single leaf from an arbitrary amount of key-value pairs.
+    /// Those pairs must all have the same leaf index.
+    fn pairs_to_leaf(pairs: Vec<(Self::Key, Self::Value)>) -> Self::Leaf;
+
     /// Maps a (MerklePath, Self::Leaf) to an opening.
     ///
     /// The length `path` is guaranteed to be equal to `DEPTH`
     fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
 
+    fn sorted_pairs_to_leaves(
+        pairs: Vec<(Self::Key, Self::Value)>,
+    ) -> PrecomputedLeaves<Self::Leaf> {
+        let mut all_leaves = PrecomputedLeaves::default();
+
+        let mut buffer: Vec<(Self::Key, Self::Value)> = Default::default();
+
+        let mut iter = pairs.into_iter().peekable();
+        while let Some((key, value)) = iter.next() {
+            let col = Self::key_to_leaf_index(&key).index.value();
+            let next_col = iter.peek().map(|(key, _)| {
+                let index = Self::key_to_leaf_index(key);
+                index.index.value()
+            });
+
+            buffer.push((key, value));
+
+            if let Some(next_col) = next_col {
+                assert!(next_col >= col);
+            }
+
+            if next_col == Some(col) {
+                // Keep going in our buffer.
+                continue;
+            }
+
+            // Whether the next pair is a different column, or non-existent, we break off.
+            let leaf_pairs = mem::take(&mut buffer);
+            let leaf = Self::pairs_to_leaf(leaf_pairs);
+            let hash = Self::hash_leaf(&leaf);
+
+            all_leaves.nodes.insert(col, leaf);
+            all_leaves.subtrees.push(SubtreeLeaf { col, hash });
+        }
+        assert_eq!(buffer.len(), 0);
+
+        all_leaves
+    }
+
     /// Builds Merkle nodes from a bottom layer of tuples of horizontal indices and their hashes,
     /// sorted by their position.
     ///
@@ -572,6 +615,24 @@ pub struct SubtreeLeaf {
     pub hash: RpoDigest,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct PrecomputedLeaves<L> {
+    /// Literal leaves to be added to the sparse Merkle tree's internal mapping.
+    pub nodes: BTreeMap<u64, L>,
+    /// "Conceptual" leaves that will be used for computations.
+    pub subtrees: Vec<SubtreeLeaf>,
+}
+
+// Derive requires `L` to impl Default, even though we don't actually need that.
+impl<L> Default for PrecomputedLeaves<L> {
+    fn default() -> Self {
+        Self {
+            nodes: Default::default(),
+            subtrees: Default::default(),
+        }
+    }
+}
+
 // TESTS
 // ================================================================================================
 #[cfg(test)]
@@ -582,9 +643,46 @@ mod test {
     use crate::{
         hash::rpo::RpoDigest,
         merkle::{smt::SubtreeLeaf, Smt, SmtLeaf, SMT_DEPTH},
-        Felt, Word, ONE,
+        Felt, Word, EMPTY_WORD, ONE,
     };
 
+    #[test]
+    fn test_sorted_pairs_to_leaves() {
+        let entries: Vec<(RpoDigest, Word)> = vec![
+            (RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]),
+            (RpoDigest::new([ONE, ONE, ONE, Felt::new(17)]), [ONE; 4]),
+            // Leaf index collision.
+            (RpoDigest::new([ONE, ONE, Felt::new(10), Felt::new(20)]), [ONE; 4]),
+            (RpoDigest::new([ONE, ONE, Felt::new(20), Felt::new(20)]), [ONE; 4]),
+            // Normal single leaf again.
+            (RpoDigest::new([ONE, ONE, ONE, Felt::new(400)]), [ONE; 4]),
+            // Empty leaf.
+            (RpoDigest::new([ONE, ONE, ONE, Felt::new(500)]), EMPTY_WORD),
+        ];
+        let mut entries_iter = entries.iter().cloned();
+        let mut next_entry = || entries_iter.next().unwrap();
+
+        let control_leaves: Vec<SmtLeaf> = vec![
+            SmtLeaf::Single(next_entry()),
+            SmtLeaf::Single(next_entry()),
+            SmtLeaf::new_multiple(vec![next_entry(), next_entry()]).unwrap(),
+            SmtLeaf::Single(next_entry()),
+            SmtLeaf::new_empty(Smt::key_to_leaf_index(&next_entry().0)),
+        ];
+
+        let control_subtree_leaves: Vec<SubtreeLeaf> = control_leaves
+            .iter()
+            .map(|leaf| {
+                let col = leaf.index().index.value();
+                let hash = leaf.hash();
+                SubtreeLeaf { col, hash }
+            })
+            .collect();
+
+        let test_subtree_leaves = Smt::sorted_pairs_to_leaves(entries).subtrees;
+        assert_eq!(control_subtree_leaves, test_subtree_leaves);
+    }
+
     #[test]
     fn test_build_subtree_from_leaves() {
         const PAIR_COUNT: u64 = u64::pow(2, 8);

src/merkle/smt/simple/mod.rs

@@ -373,4 +373,11 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
     fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath {
         (path, leaf).into()
     }
+
+    fn pairs_to_leaf(mut pairs: Vec<(LeafIndex<DEPTH>, Word)>) -> Word {
+        // SimpleSmt can't have more than one value per key.
+        assert_eq!(pairs.len(), 1);
+        let (_key, value) = pairs.pop().unwrap();
+        value
+    }
 }