store API proto AccountProof: optimize merkle node compression

CHANGELOG.md

@@ -21,6 +21,7 @@
   - Configure `NativeFaucet`, which determines the native asset used to pay fees
   - Configure the base verification fee
   - Note: fees are not yet activated, and this has no impact beyond setting these values in the block headers
+- [BREAKING] `GetAccountProofs` endpoint uses a `PartialSmt` for proofs. ([#1158](https://github.com/0xMiden/miden-node/pull/1158)).
 
 ### Fixes
 

crates/proto/Cargo.toml

@@ -25,7 +25,8 @@ thiserror        = { workspace = true }
 tonic            = { workspace = true }
 
 [dev-dependencies]
-proptest = { version = "1.7" }
+proptest          = { version = "1.7" }
+pretty_assertions = { workspace = true}
 
 [build-dependencies]
 anyhow                 = { workspace = true }

crates/proto/src/domain/merkle.rs

@@ -1,9 +1,14 @@
+use std::collections::{HashMap, HashSet};
+
 use miden_objects::Word;
 use miden_objects::crypto::merkle::{
     Forest,
+    InnerNode,
     LeafIndex,
     MerklePath,
     MmrDelta,
+    NodeIndex,
+    PartialSmt,
     SmtLeaf,
     SmtProof,
     SparseMerklePath,
@@ -206,3 +211,226 @@ impl From<SmtProof> for proto::primitives::SmtOpening {
         }
     }
 }
+
+// NODE INDEX
+// ------------------------------------------------------------------------------------------------
+impl From<NodeIndex> for proto::primitives::NodeIndex {
+    fn from(value: NodeIndex) -> Self {
+        proto::primitives::NodeIndex {
+            depth: value.depth() as u32,
+            value: value.value(),
+        }
+    }
+}
+impl TryFrom<proto::primitives::NodeIndex> for NodeIndex {
+    type Error = ConversionError;
+    fn try_from(index: proto::primitives::NodeIndex) -> Result<Self, Self::Error> {
+        let depth = u8::try_from(index.depth)?;
+        let value = index.value;
+        Ok(NodeIndex::new(depth, value)?)
+    }
+}
+
+// PARTIAL SMT
+// ------------------------------------------------------------------------------------------------
+
+impl TryFrom<proto::primitives::PartialSmt> for PartialSmt {
+    type Error = ConversionError;
+    fn try_from(value: proto::primitives::PartialSmt) -> Result<Self, Self::Error> {
+        let proto::primitives::PartialSmt { root, leaves, nodes } = value;
+        let root = root
+            .as_ref()
+            .ok_or(proto::primitives::PartialSmt::missing_field(stringify!(root)))?
+            .try_into()?;
+        // TODO ensure `!leaves.is_empty()`
+
+        // Convert other proto primitives to crypto types
+        let leaves = Result::<Vec<SmtLeaf>, _>::from_iter(try_convert(leaves))?;
+        let mut inner =
+            Result::<HashMap<NodeIndex, Word>, _>::from_iter(nodes.into_iter().map(|inner| {
+                let node_index = NodeIndex::try_from(
+                    inner
+                        .index
+                        .ok_or(proto::primitives::NodeIndex::missing_field(stringify!(index)))?,
+                )?;
+                let digest = Word::try_from(
+                    inner
+                        .digest
+                        .ok_or(proto::primitives::Digest::missing_field(stringify!(digest)))?,
+                )?;
+                Ok::<_, Self::Error>((node_index, digest))
+            }))?;
+
+        let leaf_indices =
+            HashSet::<NodeIndex>::from_iter(leaves.iter().map(|leaf| leaf.index().into()));
+
+        // Must contain the leaves too
+        inner.extend(leaves.iter().map(|leaf| (leaf.index().into(), leaf.hash())));
+
+        // Start constructing the partial SMT
+        //
+        // Construct a `MerklePath` per leaf by transcending from leaf digest down to depth 0.
+        // Then verify the merkle proof holds consistency and completeness checks and all
+        // required sibling nodes are present to deeduct required intermediate nodes.
+        let mut partial = PartialSmt::new();
+        for leaf in leaves {
+            // Construct the merkle path:
+            let leaf_node_index: NodeIndex = leaf.index().into();
+            let mut current = leaf_node_index.clone();
+            let mut siblings = Vec::new();
+
+            // If we ever try to trancend beyond this depth level, something is wrong and
+            // we must stop decoding.
+            let max_depth = leaf_node_index.depth();
+            // root:                00
+            //                    /    \
+            //                 10         11
+            //                /  \       /  \
+            //              20   21     22   23
+            //              / \   / \  / \   / \
+            //              leaves ...       x  y
+            // Iterate from the leaf up to the root (exclusive)
+            // We start by picking the sibling of `x`, `y`, our starting point and
+            // then moving towards the root `0`. By definition siblings have the same parent.
+            loop {
+                let sibling_idx = current.sibling();
+                // TODO FIXME for a leaf we get another leaf, we need to ensure those are part of
+                // the inner set or contained in the inner HashMap
+                let sibling_digest = if let Some(sibling_digest) = inner.get(&sibling_idx) {
+                    // Previous round already calculated the entry or it was given explicitly
+                    *sibling_digest
+                } else {
+                    // The entry does not exist, so we need to lazily follow the missing nodes and
+                    // calculate recursively.
+
+                    // DFS, build the subtree recursively, starting from the current sibling
+                    let mut stack = Vec::<NodeIndex>::new();
+                    stack.push(sibling_idx.clone());
+                    loop {
+                        let Some(idx) = stack.pop() else {
+                            unreachable!(
+                                "Must be an error, we must have nodes to resolve all questions, otherwise construction is borked"
+                            )
+                        };
+                        if let Some(node_digest) = inner.get(&idx) {
+                            if stack.is_empty() && idx == sibling_idx {
+                                // we emptied the stack which means the current one is our desired
+                                // starting point
+                                break *node_digest;
+                            }
+                            // if the digest exists, we don't need to recurse
+                            continue;
+                        }
+                        debug_assert!(
+                            !leaf_indices.contains(&idx),
+                            "For every relevant leaf, we must have the relevant value"
+                        );
+                        let left = idx.left_child();
+                        let right = idx.right_child();
+                        if max_depth < left.depth() || max_depth < right.depth() {
+                            // TODO might happen in case of a missing node, so we must handle this
+                            // gracefully
+                            unreachable!("graceful!")
+                        }
+                        // proceed if the inner nodes are unknown
+                        if !inner.contains_key(&left) {
+                            stack.push(left);
+                        }
+                        if !inner.contains_key(&right) {
+                            stack.push(right);
+                        }
+                        // left and right exist, we can derive the digest for `idx`
+                        if let Some(&left) = inner.get(&left)
+                            && let Some(&right) = inner.get(&right)
+                        {
+                            let node = InnerNode { left, right };
+                            let node_digest = node.hash();
+
+                            if stack.is_empty() && idx == sibling_idx {
+                                // we emptied the stack which means the current one is our desired
+                                // starting point
+                                break node_digest;
+                            }
+                            inner.insert(idx, node_digest);
+                        }
+                    }
+                };
+                siblings.push(sibling_digest);
+
+                // Move up to the parent level, and repeat
+                current = current.parent();
+                if current.depth() == 0 {
+                    break;
+                }
+            }
+
+            let path = MerklePath::new(siblings);
+            path.verify(leaf_node_index.value(), leaf.hash(), &root).expect("It's fine");
+            partial.add_path(leaf, path);
+        }
+        assert_eq!(partial.root(), root); // FIXME make error
+        Ok(partial)
+    }
+}
+
+impl From<PartialSmt> for proto::primitives::PartialSmt {
+    fn from(partial: PartialSmt) -> Self {
+        // Find all leaf digests, we need to include those, they are POIs
+        let mut leaves = Vec::new();
+        for (key, value) in partial.entries() {
+            let leaf = partial.get_leaf(key).unwrap();
+            leaves.push(crate::generated::primitives::SmtLeaf::from(leaf));
+        }
+
+        // Now collect the minimal set of internal nodes to be able to recalc the intermediate nodes
+        // forming a partial smt
+        let mut retained = HashMap::<NodeIndex, Word>::new();
+        for (idx, node) in partial.inner_node_indices() {
+            // if neither of the child keys are tracked, we cannot re-calc the inner node digest
+            // on-the-fly and hence need to add the node to the set to be transferred
+            if partial.get_value(node.left).is_err() || partial.get_value(node.left).is_err() {
+                retained.insert(idx, node.hash());
+                continue;
+            }
+        }
+        let nodes = Vec::from_iter(retained.into_iter().map(|(index, digest)| {
+            crate::generated::primitives::InnerNode {
+                index: Some(crate::generated::primitives::NodeIndex::from(index)),
+                digest: Some(crate::generated::primitives::Digest::from(digest)),
+            }
+        }));
+        let root = Some(partial.root().into());
+        // Remember: nodes and leaves as mutually exclusive
+        Self { root, nodes, leaves }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use miden_objects::crypto::merkle::{PartialSmt, Smt};
+    use pretty_assertions::assert_eq;
+
+    use super::*;
+    #[test]
+    fn partial_smt_roundtrip() {
+        let mut x = Smt::new();
+
+        x.insert(Word::from([1_u32, 2, 3, 4]), Word::from([5_u32, 6, 7, 8]));
+        x.insert(Word::from([10_u32, 11, 12, 13]), Word::from([14_u32, 15, 16, 17]));
+        x.insert(Word::from([0x00_u32, 0xFF, 0xFF, 0xFF]), Word::from([0x00_u32; 4]));
+        x.insert(Word::from([0xAA_u32, 0xFF, 0xFF, 0xFF]), Word::from([0xAA_u32; 4]));
+        x.insert(Word::from([0xBB_u32, 0xFF, 0xFF, 0xFF]), Word::from([0xBB_u32; 4]));
+        x.insert(Word::from([0xCC_u32, 0xFF, 0xFF, 0xFF]), Word::from([0xCC_u32; 4]));
+
+        let proof = x.open(&Word::from([10_u32, 11, 12, 13]));
+
+        let mut orig = PartialSmt::new();
+        orig.add_proof(proof);
+        let orig = orig;
+
+        let proto = proto::primitives::PartialSmt::from(orig.clone());
+        let recovered = PartialSmt::try_from(proto).unwrap();
+
+        assert_eq!(orig, recovered);
+    }
+}

crates/proto/src/generated/primitives.rs

@@ -1,4 +1,39 @@
 // This file is @generated by prost-build.
+/// Representation of a partial sparse merkle tree.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct PartialSmt {
+    /// The sparse merkle tree root.
+    #[prost(message, optional, tag = "1")]
+    pub root: ::core::option::Option<Digest>,
+    /// Set of leaves of the merkle tree.
+    #[prost(message, repeated, tag = "2")]
+    pub leaves: ::prost::alloc::vec::Vec<SmtLeaf>,
+    /// Unique set of inner merkle tree digest, all belonging to at least one
+    /// merkle path of a given leave. Note that we skip all inner nodes that do
+    /// have two children, since we can recalculate them on-the-fly.
+    #[prost(message, repeated, tag = "3")]
+    pub nodes: ::prost::alloc::vec::Vec<InnerNode>,
+}
+/// Node index representation.
+#[derive(Clone, Copy, PartialEq, ::prost::Message)]
+pub struct NodeIndex {
+    /// The depth of the index, starting from 0 as root.
+    #[prost(uint32, tag = "1")]
+    pub depth: u32,
+    /// The index within a certain tree depth, left-most being zero.
+    #[prost(uint64, tag = "2")]
+    pub value: u64,
+}
+/// Inner node of a sparse merkle tree
+#[derive(Clone, Copy, PartialEq, ::prost::Message)]
+pub struct InnerNode {
+    /// The position of the inner node within the tree.
+    #[prost(message, optional, tag = "1")]
+    pub index: ::core::option::Option<NodeIndex>,
+    /// The digest of the subtree down to the root.
+    #[prost(message, optional, tag = "2")]
+    pub digest: ::core::option::Option<Digest>,
+}
 /// Represents a single SMT leaf entry.
 #[derive(Clone, Copy, PartialEq, ::prost::Message)]
 pub struct SmtLeafEntry {

crates/proto/src/generated/rpc_store.rs

@@ -99,23 +99,25 @@ pub mod account_proofs {
             /// the current one.
             #[prost(bytes = "vec", optional, tag = "3")]
             pub account_code: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
-            /// Storage slots information for this account
-            #[prost(message, repeated, tag = "4")]
-            pub storage_maps: ::prost::alloc::vec::Vec<
-                account_state_header::StorageSlotMapProof,
+            /// A sparse merkle tree per storage slot, including all relevant merkle proofs for storage entries.
+            #[prost(message, repeated, tag = "5")]
+            pub partial_storage_smts: ::prost::alloc::vec::Vec<
+                account_state_header::StorageSlotMapPartialSmt,
             >,
         }
         /// Nested message and enum types in `AccountStateHeader`.
         pub mod account_state_header {
             /// Represents a single storage slot with the requested keys and their respective values.
             #[derive(Clone, PartialEq, ::prost::Message)]
-            pub struct StorageSlotMapProof {
+            pub struct StorageSlotMapPartialSmt {
                 /// The storage slot index (\[0..255\]).
                 #[prost(uint32, tag = "1")]
                 pub storage_slot: u32,
-                /// Merkle proof of the map value
-                #[prost(bytes = "vec", tag = "2")]
-                pub smt_proof: ::prost::alloc::vec::Vec<u8>,
+                /// Merkle proofs of the map value as partial sparse merkle tree for compression.
+                #[prost(message, optional, tag = "2")]
+                pub partial_smt: ::core::option::Option<
+                    super::super::super::super::primitives::PartialSmt,
+                >,
             }
         }
     }