From 8283ec5c3066f7fbcd77e47ab71e78606af8aa9a Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Wed, 13 Nov 2024 12:29:05 -0700
Subject: [PATCH] smt: add pairs_to_leaf() to trait

---
 src/merkle/smt/full/mod.rs   | 24 ++++++++++++++++++++++++
 src/merkle/smt/mod.rs        |  4 ++++
 src/merkle/smt/simple/mod.rs | 12 +++++++++++-
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs
index 226a8b1..1e2c574 100644
--- a/src/merkle/smt/full/mod.rs
+++ b/src/merkle/smt/full/mod.rs
@@ -344,6 +344,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 {
diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index 056c221..03d9d45 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -338,6 +338,10 @@ 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`
diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs
index 6229ac2..04476a0 100644
--- a/src/merkle/smt/simple/mod.rs
+++ b/src/merkle/smt/simple/mod.rs
@@ -1,4 +1,7 @@
-use alloc::collections::{BTreeMap, BTreeSet};
+use alloc::{
+    collections::{BTreeMap, BTreeSet},
+    vec::Vec,
+};
 
 use super::{
     super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,
@@ -370,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
+    }
 }