test_mock_proof_public_data_matches_witnesses

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

@@ -537,21 +537,7 @@ fn compute_state_update_namespace<S: MerkleProofSpec>(
     let mut finished = session.finish(accesses)?;
     if write_witness {
         let nomt_witness = finished.take_witness().expect("Witness cannot be missing");
-        let nomt::Witness {
-            path_proofs,
-            operations: nomt::WitnessedOperations { .. },
-        } = nomt_witness;
-        // Note, we discard `p.path`, but maybe there's a way to use to have more efficient verification?
-        let mut path_proofs_inner = path_proofs.into_iter().map(|p| p.inner).collect::<Vec<_>>();
-
-        // Sort them as required by
-        // Note that the path proofs produced within a crate::witness::Witness are not guaranteed to be ordered,
-        // so the input should be sorted lexicographically by the terminal path prior to calling this function.
-        // https://github.com/thrumdev/nomt/issues/904
-        path_proofs_inner.sort_by(|a, b| a.terminal.path().cmp(b.terminal.path()));
-
-        let multi_proof = MultiProof::from_path_proofs(path_proofs_inner);
-        witness.add_hint(&multi_proof);
+        witness.add_hint(&nomt_witness);
     }
     Ok(finished)
 }

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

@@ -1,9 +1,11 @@
 //! ZK Verifier part of the NOMT based Storage implementation
+use std::collections::BTreeMap;
 use std::marker::PhantomData;
 
 use nomt_core::hasher::BinaryHasher;
-use nomt_core::proof::MultiProof;
+use nomt_core::proof::{MultiProof, PathUpdate};
 use nomt_core::trie::{KeyPath, LeafData, Node, ValueHash};
+use nomt_core::witness::Witness as NomtWitness;
 #[cfg(all(feature = "test-utils", feature = "native"))]
 use sov_rollup_interface::common::SlotNumber;
 use sov_rollup_interface::reexports::digest::Digest;
@@ -41,69 +43,173 @@ impl<S: MerkleProofSpec> NomtVerifierStorage<S> {
             ordered_reads: state_reads,
             ordered_writes: state_writes,
         } = state_accesses;
+        let mut expected_reads = BTreeMap::new();
+        for (key, value) in state_reads {
+            let key_hash: KeyPath = S::Hasher::digest(key.as_ref()).into();
+            let value_hash = value.map(|node_leaf| {
+                S::Hasher::digest(node_leaf.combine_val_hash_and_size()).into()
+            });
+            if expected_reads.insert(key_hash, value_hash).is_some() {
+                anyhow::bail!("Duplicate key read in state accesses: {:?}", key);
+            }
+        }
+
+        let mut expected_writes = BTreeMap::new();
+        for (key, value) in state_writes {
+            let key_hash: KeyPath = S::Hasher::digest(key.as_ref()).into();
+            let value_hash = value.map(|slot_value| {
+                // Authenticated write is hash of a combination of size and original value hash.
+                S::Hasher::digest(slot_value.combine_val_hash_and_size::<S::Hasher>()).into()
+            });
+            if expected_writes.insert(key_hash, value_hash).is_some() {
+                anyhow::bail!("Duplicate key write in state accesses: {:?}", key);
+            }
+        }
 
-        let multi_proof: MultiProof = array_witness.get_hint();
+        let nomt_witness: NomtWitness = array_witness.get_hint();
+        let mut path_proofs_inner = nomt_witness
+            .path_proofs
+            .iter()
+            .map(|path| path.inner.clone())
+            .collect::<Vec<_>>();
+        path_proofs_inner.sort_by(|a, b| a.terminal.path().cmp(b.terminal.path()));
+        let multi_proof = MultiProof::from_path_proofs(path_proofs_inner);
         let verified_multi_proof = nomt_core::proof::verify_multi_proof::<BinaryHasher<S::Hasher>>(
             &multi_proof,
             prev_root,
         )
         .map_err(|e| anyhow::anyhow!("Failed to verify multi proof: {:?}", e))?;
 
-        for (key, value) in state_reads {
-            let key_hash: KeyPath = S::Hasher::digest(key.as_ref()).into();
+        for (key, value) in &expected_reads {
             match value {
                 None => {
                     if !verified_multi_proof
-                        .confirm_nonexistence(&key_hash)
+                        .confirm_nonexistence(key)
                         .map_err(|e| anyhow::anyhow!("Failed to confirm non-existence: {:?}", e))?
                     {
-                        anyhow::bail!("Failed to verify non-existence of key: {:?}", key);
+                        anyhow::bail!("Failed to verify non-existence of key");
                     }
                 }
-                Some(node_leaf) => {
-                    let authenticated_write = node_leaf.combine_val_hash_and_size();
-                    let value_hash = S::Hasher::digest(&authenticated_write).into();
+                Some(value_hash) => {
                     let leaf = LeafData {
-                        key_path: key_hash,
-                        value_hash,
+                        key_path: *key,
+                        value_hash: *value_hash,
                     };
                     if !verified_multi_proof
                         .confirm_value(&leaf)
                         .map_err(|e| anyhow::anyhow!("Failed to confirm value: {:?}", e))?
                     {
-                        anyhow::bail!("Failed to verify inclusion of key: {:?}", key);
+                        anyhow::bail!("Failed to verify inclusion of key");
                     }
                 }
             }
         }
 
-        let mut updates = state_writes
+        let mut witnessed_writes = BTreeMap::new();
+        for write in nomt_witness.operations.writes {
+            let Some(expected_value) = expected_writes.remove(&write.key) else {
+                anyhow::bail!("Unexpected or duplicate NOMT write witness entry");
+            };
+            if write.value != expected_value {
+                anyhow::bail!(
+                    "Mismatched NOMT write witness value for key {:?}: witness={:?}, expected={:?}",
+                    write.key,
+                    write.value,
+                    expected_value
+                );
+            }
+            if witnessed_writes
+                .insert(write.key, expected_value)
+                .is_some()
+            {
+                anyhow::bail!("Duplicate NOMT write witness entry");
+            };
+        }
+
+        if !expected_writes.is_empty() {
+            anyhow::bail!("Missing NOMT write witness entries");
+        }
+
+        let mut verified_paths = nomt_witness
+            .path_proofs
             .into_iter()
-            .map(|(key, value)| {
-                (
-                    S::Hasher::digest(key.as_ref()).into(),
-                    value.map(|slot_value| {
-                        // Authenticated write is hash of a combination of size and orignal value hash.
-                        S::Hasher::digest(slot_value.combine_val_hash_and_size::<S::Hasher>())
-                            .into()
-                    }),
-                )
+            .map(|path| {
+                let verified = path
+                    .inner
+                    .verify::<BinaryHasher<S::Hasher>>(path.path.path(), prev_root)
+                    .map_err(|e| anyhow::anyhow!("Failed to verify path proof: {:?}", e))?;
+                Ok((
+                    verified,
+                    path.path.depth() as usize,
+                    truncate_key_path(path.path.raw_path(), path.path.depth() as usize),
+                    Vec::<(KeyPath, Option<ValueHash>)>::new(),
+                ))
             })
-            .collect::<Vec<(KeyPath, Option<ValueHash>)>>();
+            .collect::<anyhow::Result<Vec<_>>>()?;
+        verified_paths.sort_by(|a, b| a.0.path().cmp(b.0.path()));
+
+        let mut path_index_by_prefix = BTreeMap::new();
+        let mut path_depths_desc = Vec::new();
+        for (index, (_, depth, prefix, _)) in verified_paths.iter().enumerate() {
+            if path_index_by_prefix
+                .insert((*depth, *prefix), index)
+                .is_some()
+            {
+                anyhow::bail!("Duplicate NOMT path proof prefix");
+            }
+            path_depths_desc.push(*depth);
+        }
+        path_depths_desc.sort_unstable();
+        path_depths_desc.dedup();
+        path_depths_desc.reverse();
+
+        for (key, value) in witnessed_writes {
+            let matching_index = path_depths_desc
+                .iter()
+                .find_map(|depth| {
+                    path_index_by_prefix
+                        .get(&(*depth, truncate_key_path(key, *depth)))
+                        .copied()
+                })
+                .ok_or_else(|| anyhow::anyhow!("No NOMT path proof covers write key {:?}", key))?;
+
+            verified_paths[matching_index].3.push((key, value));
+        }
 
-        // Sort them by key hash, as required by [`nomt_core::proof::verify_multi_proof_update`]
-        updates.sort_by(|a, b| a.0.cmp(&b.0));
+        let mut updates = Vec::new();
+        for (verified, _, _, writes) in verified_paths {
+            if !writes.is_empty() {
+                updates.push(PathUpdate {
+                    inner: verified,
+                    ops: writes,
+                });
+            }
+        }
 
-        nomt_core::proof::verify_multi_proof_update::<BinaryHasher<S::Hasher>>(
-            &verified_multi_proof,
-            updates,
-        )
-        .map_err(|e| anyhow::anyhow!("Failed to verify update: {:?}", e))
-        // Note: we don't check exhaustion of the proof
-        // because it does not impact the correctness of the guest, only performance.
+        nomt_core::proof::verify_update::<BinaryHasher<S::Hasher>>(prev_root, &updates)
+            .map_err(|e| anyhow::anyhow!("Failed to verify update: {:?}", e))
     }
 }
 
+fn truncate_key_path(mut key: KeyPath, depth: usize) -> KeyPath {
+    debug_assert!(depth <= 256);
+
+    let full_bytes = depth / 8;
+    let partial_bits = depth % 8;
+
+    if full_bytes < key.len() {
+        if partial_bits == 0 {
+            key[full_bytes..].fill(0);
+        } else {
+            let keep_mask = u8::MAX << (8 - partial_bits);
+            key[full_bytes] &= keep_mask;
+            key[full_bytes + 1..].fill(0);
+        }
+    }
+
+    key
+}
+
 impl<S: MerkleProofSpec> Storage for NomtVerifierStorage<S> {
     type Hasher = S::Hasher;
     type Witness = S::Witness;

examples/demo-rollup/tests/prover/sp1_cpu_prover.rs

@@ -20,7 +20,7 @@ type ProofInput = StateTransitionWitnessWithAddress<
 >;
 
 #[tokio::test(flavor = "multi_thread")]
-#[ignore = "This test is used to generate data for testing the aggregate proof circuit and should be enabled only when needed."]
+//#[ignore = "This test is used to generate data for testing the aggregate proof circuit and should be enabled only when needed."]
 async fn test_save_proofs() {
     let (host, code_commitment) = TestHost::new(true).await;
     let proof_data = generate_proofs(&host).await;
@@ -55,7 +55,8 @@ async fn generate_proofs(host: &TestHost) -> Vec<BlockHeaderWithProof<MockDaSpec
     let prover_address = <DefaultSpec as Spec>::Address::try_from([0u8; 28].as_ref()).unwrap();
     let mut proofs = Vec::new();
 
-    for witness in witnesses {
+    for (i, witness) in witnesses.into_iter().enumerate() {
+        println!("XX {i}");
         let da_block_header = witness.da_block_header.clone();
 
         let data: ProofInput = StateTransitionWitnessWithAddress {
@@ -73,6 +74,30 @@ async fn generate_proofs(host: &TestHost) -> Vec<BlockHeaderWithProof<MockDaSpec
     proofs
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_proof_public_data_matches_witnesses() {
+    let (host, _) = TestHost::new(false).await;
+    let (_genesis_state_root, witnesses) = super::generate_witnesses().await;
+    let prover_address = <DefaultSpec as Spec>::Address::try_from([0u8; 28].as_ref()).unwrap();
+
+    for witness in witnesses {
+        let expected_initial_state_root = witness.initial_state_root.clone();
+        let expected_final_state_root = witness.final_state_root.clone();
+        let expected_slot_hash = witness.da_block_header.hash();
+
+        let public_data = host
+            .run_mock_public_data(ProofInput {
+                stf_witness: witness,
+                prover_address,
+            })
+            .await;
+
+        assert_eq!(public_data.initial_state_root, expected_initial_state_root);
+        assert_eq!(public_data.final_state_root, expected_final_state_root);
+        assert_eq!(public_data.slot_hash, expected_slot_hash);
+    }
+}
+
 // The SP1 prover manages its own Tokio runtime, which conflicts with the `tokio::test` runtime.
 // To avoid this, all blocking work must be executed inside `tokio::task::spawn_blocking`.
 struct TestHost {
@@ -129,6 +154,25 @@ impl TestHost {
             .unwrap()
         }
     }
+
+    async fn run_mock_public_data(
+        &self,
+        data: ProofInput,
+    ) -> StateTransitionPublicData<<DefaultSpec as Spec>::Address, MockDaSpec, ProofStateRoot> {
+        tokio::task::spawn_blocking(move || {
+            let mock_host = MockSp1Prover::new(*sp1::SP1_GUEST_MOCK_ELF)
+                .expect("MockSp1Prover should be created successfully");
+
+            let proof = mock_host
+                .add_hint_and_run(&data)
+                .expect("Mock prover should run successfully");
+
+            bincode::deserialize(proof.public_values.as_slice())
+                .expect("Mock proof public values should deserialize")
+        })
+        .await
+        .unwrap()
+    }
 }
 
 #[allow(dead_code)]