Use update proofs

Cargo.toml

@@ -54,8 +54,7 @@ rust-version = "1.66"
 
 [workspace.dependencies]
 # Dependencies maintained by sovereign
-# TODO: replace by release number once available on crates.io: tracking issue https://github.com/Sovereign-Labs/sovereign-sdk/issues/632
-jmt = { git = "https://github.com/penumbra-zone/jmt", commit = "46b4b00" }
+jmt = "0.7.0"
 
 # External dependencies
 async-trait = "0.1.71"
@@ -112,4 +111,5 @@ secp256k1 = { version = "0.27.0", default-features = false, features = ["global-
 [patch.crates-io]
 # See reth: https://github.com/paradigmxyz/reth/blob/main/Cargo.toml#L79
 revm = { git = "https://github.com/bluealloy/revm/", branch = "release/v25" }
-revm-primitives = { git = "https://github.com/bluealloy/revm/", branch = "release/v25" }
+revm-primitives = { git = "https://github.com/bluealloy/revm/", branch = "release/v25" }
+jmt = { path = "../penumbra-jmt/jmt" }

module-system/sov-state/src/lib.rs

@@ -5,9 +5,6 @@ mod map;
 #[cfg(feature = "native")]
 mod prover_storage;
 
-#[cfg(feature = "native")]
-mod tree_db;
-
 mod scratchpad;
 
 pub mod storage;
@@ -37,7 +34,7 @@ pub use storage::Storage;
 use utils::AlignedVec;
 pub use value::StateValue;
 
-pub use crate::witness::{ArrayWitness, TreeWitnessReader, Witness};
+pub use crate::witness::{ArrayWitness, Witness};
 
 // A prefix prepended to each key before insertion and retrieval from the storage.
 // All the collection types in this crate are backed by the same storage instance, this means that insertions of the same key

module-system/sov-state/src/prover_storage.rs

@@ -10,7 +10,6 @@ use sov_db::state_db::StateDB;
 use crate::config::Config;
 use crate::internal_cache::OrderedReadsAndWrites;
 use crate::storage::{NativeStorage, StorageKey, StorageProof, StorageValue};
-use crate::tree_db::TreeReadLogger;
 use crate::witness::Witness;
 use crate::{MerkleProofSpec, Storage};
 
@@ -85,19 +84,14 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
         witness: &Self::Witness,
     ) -> Result<[u8; 32], anyhow::Error> {
         let latest_version = self.db.get_next_version() - 1;
-        witness.add_hint(latest_version);
 
-        let read_logger = TreeReadLogger::with_db_and_witness(self.db.clone(), witness);
-        let untracked_jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&self.db);
+        let jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&self.db);
 
         // Handle empty untracked_jmt
-        if untracked_jmt
-            .get_root_hash_option(latest_version)?
-            .is_none()
-        {
+        if jmt.get_root_hash_option(latest_version)?.is_none() {
             assert_eq!(latest_version, 0);
             let empty_batch = Vec::default().into_iter();
-            let (_, tree_update) = untracked_jmt
+            let (_, tree_update) = jmt
                 .put_value_set(empty_batch, latest_version)
                 .expect("JMT update must succeed");
 
@@ -110,14 +104,13 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
         for (key, read_value) in state_accesses.ordered_reads {
             let key_hash = KeyHash::with::<S::Hasher>(key.key.as_ref());
             // TODO: Switch to the batch read API once it becomes available
-            let (result, proof) = untracked_jmt.get_with_proof(key_hash, latest_version)?;
+            let (result, proof) = jmt.get_with_proof(key_hash, latest_version)?;
             if result.as_ref() != read_value.as_ref().map(|f| f.value.as_ref()) {
                 anyhow::bail!("Bug! Incorrect value read from jmt");
             }
             witness.add_hint(proof);
         }
 
-        let tracked_jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&read_logger);
         // Compute the jmt update from the write batch
         let batch = state_accesses
             .ordered_writes
@@ -135,13 +128,17 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
 
         let next_version = self.db.get_next_version();
 
-        let (new_root, tree_update) = tracked_jmt
-            .put_value_set(batch, next_version)
+        let (new_root, update_proof, tree_update) = jmt
+            .put_value_set_with_proof(batch, next_version)
             .expect("JMT update must succeed");
 
         self.db
             .write_node_batch(&tree_update.node_batch)
             .expect("db write must succeed");
+
+        witness.add_hint(update_proof);
+        witness.add_hint(&new_root.0);
+
         self.db.inc_next_version();
         Ok(new_root.0)
     }

module-system/sov-state/src/witness.rs

@@ -2,7 +2,6 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::Mutex;
 
 use borsh::{BorshDeserialize, BorshSerialize};
-use jmt::storage::TreeReader;
 use serde::{Deserialize, Serialize};
 
 // TODO: Refactor witness trait so it only require Serialize / Deserialize
@@ -13,42 +12,6 @@ pub trait Witness: Default + Serialize {
     fn merge(&self, rhs: &Self);
 }
 
-#[derive(Debug)]
-pub struct TreeWitnessReader<'a, T: Witness>(&'a T);
-
-impl<'a, T: Witness> TreeWitnessReader<'a, T> {
-    pub fn new(witness: &'a T) -> Self {
-        Self(witness)
-    }
-}
-
-impl<'a, T: Witness> TreeReader for TreeWitnessReader<'a, T> {
-    fn get_node_option(
-        &self,
-        _node_key: &jmt::storage::NodeKey,
-    ) -> anyhow::Result<Option<jmt::storage::Node>> {
-        let serialized_node_opt: Option<Vec<u8>> = self.0.get_hint();
-        match serialized_node_opt {
-            Some(val) => Ok(Some(jmt::storage::Node::deserialize_reader(&mut &val[..])?)),
-            None => Ok(None),
-        }
-    }
-
-    fn get_value_option(
-        &self,
-        _max_version: jmt::Version,
-        _key_hash: jmt::KeyHash,
-    ) -> anyhow::Result<Option<jmt::OwnedValue>> {
-        Ok(self.0.get_hint())
-    }
-
-    fn get_rightmost_leaf(
-        &self,
-    ) -> anyhow::Result<Option<(jmt::storage::NodeKey, jmt::storage::LeafNode)>> {
-        unimplemented!()
-    }
-}
-
 #[derive(Default, Debug, Serialize, Deserialize)]
 pub struct ArrayWitness {
     next_idx: AtomicUsize,

module-system/sov-state/src/zk_storage.rs

@@ -1,13 +1,13 @@
 use std::marker::PhantomData;
 use std::sync::Arc;
 
-use jmt::{JellyfishMerkleTree, KeyHash, Version};
+use jmt::{KeyHash, RootHash};
 #[cfg(all(target_os = "zkvm", feature = "bench"))]
 use zk_cycle_macros::cycle_tracker;
 
 use crate::internal_cache::OrderedReadsAndWrites;
 use crate::storage::{StorageKey, StorageProof, StorageValue};
-use crate::witness::{TreeWitnessReader, Witness};
+use crate::witness::Witness;
 use crate::{MerkleProofSpec, Storage};
 
 #[cfg(all(target_os = "zkvm", feature = "bench"))]
@@ -61,9 +61,6 @@ impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
         state_accesses: OrderedReadsAndWrites,
         witness: &Self::Witness,
     ) -> Result<[u8; 32], anyhow::Error> {
-        let latest_version: Version = witness.get_hint();
-        let reader = TreeWitnessReader::new(witness);
-
         // For each value that's been read from the tree, verify the provided smt proof
         for (key, read_value) in state_accesses.ordered_reads {
             let key_hash = KeyHash::with::<S::Hasher>(key.key.as_ref());
@@ -79,6 +76,9 @@ impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
             }
         }
 
+        let update_proof: jmt::proof::UpdateMerkleProof<S::Hasher> = witness.get_hint();
+        let new_root = RootHash(witness.get_hint());
+
         // Compute the jmt update from the write batch
         let batch = state_accesses
             .ordered_writes
@@ -89,16 +89,15 @@ impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
                     key_hash,
                     value.map(|v| Arc::try_unwrap(v.value).unwrap_or_else(|arc| (*arc).clone())),
                 )
-            });
+            })
+            .collect::<Vec<_>>();
 
-        let next_version = latest_version + 1;
         // TODO: Make updates verifiable. Currently, writes don't verify that the provided siblings existed in the old tree
         // because the TreeReader is trusted
-        let jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&reader);
+        update_proof
+            .verify_update(RootHash(self.prev_state_root), new_root, batch)
+            .expect("Update proof was invalid! The prover was malicious");
 
-        let (new_root, _tree_update) = jmt
-            .put_value_set(batch, next_version)
-            .expect("JMT update must succeed");
         Ok(new_root.0)
     }