From 3157c55de64d5928586b7f9f19530dd4c42ed12d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Kr=C3=BCger?= <ole.kruger@trili.tech>
Date: Thu, 1 May 2025 13:01:48 +0100
Subject: [PATCH] RISC-V: Consolidate proof generation

---
 src/riscv/lib/src/pvm/common.rs   | 18 ++++++++++++++++++
 src/riscv/lib/src/pvm/node_pvm.rs | 12 ++----------
 src/riscv/lib/src/stepper/pvm.rs  | 17 ++++-------------
 3 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/src/riscv/lib/src/pvm/common.rs b/src/riscv/lib/src/pvm/common.rs
index 7778ac31625e..24bdb74843cc 100644
--- a/src/riscv/lib/src/pvm/common.rs
+++ b/src/riscv/lib/src/pvm/common.rs
@@ -35,10 +35,12 @@ use crate::state_backend::FnManagerIdent;
 use crate::state_backend::ManagerBase;
 use crate::state_backend::ProofLayout;
 use crate::state_backend::ProofTree;
+use crate::state_backend::Ref;
 use crate::state_backend::owned_backend::Owned;
 use crate::state_backend::proof_backend::ProofGen;
 use crate::state_backend::proof_backend::ProofWrapper;
 use crate::state_backend::proof_backend::proof::MerkleProof;
+use crate::state_backend::proof_backend::proof::Proof;
 use crate::state_backend::verify_backend::Verifier;
 use crate::storage::Hash;
 use crate::storage::HashError;
@@ -507,6 +509,22 @@ impl<MC: MemoryConfig, CL: CacheLayouts, B: Block<MC, Owned>> Pvm<MC, CL, B, Own
     }
 }
 
+impl<'a, MC: MemoryConfig, CL: CacheLayouts, B: Block<MC, ProofGen<Ref<'a, Owned>>>>
+    Pvm<MC, CL, B, ProofGen<Ref<'a, Owned>>>
+{
+    /// Produce a proof.
+    pub(crate) fn to_proof(&self) -> Result<Proof, HashError> {
+        let refs = self.struct_ref::<FnManagerIdent>();
+        let merkle_proof = PvmLayout::<MC, CL>::to_merkle_tree(refs)?.to_merkle_proof()?;
+
+        let refs = self.struct_ref::<FnManagerIdent>();
+        let final_hash = PvmLayout::<MC, CL>::state_hash(refs)?;
+        let proof = Proof::new(merkle_proof, final_hash);
+
+        Ok(proof)
+    }
+}
+
 impl<MC: MemoryConfig, CL: CacheLayouts, B: Block<MC, Verifier>> Pvm<MC, CL, B, Verifier> {
     /// Construct a PVM state from a Merkle proof.
     pub fn from_proof(proof: &MerkleProof, block_builder: B::BlockBuilder) -> Option<Self> {
diff --git a/src/riscv/lib/src/pvm/node_pvm.rs b/src/riscv/lib/src/pvm/node_pvm.rs
index cbcea6bd7a76..bdc102f59422 100644
--- a/src/riscv/lib/src/pvm/node_pvm.rs
+++ b/src/riscv/lib/src/pvm/node_pvm.rs
@@ -21,7 +21,6 @@ use crate::pvm::common::PvmStatus;
 use crate::state::NewState;
 use crate::state_backend;
 use crate::state_backend::AllocatedOf;
-use crate::state_backend::CommitmentLayout;
 use crate::state_backend::FnManagerIdent;
 use crate::state_backend::ProofLayout;
 use crate::state_backend::ProofTree;
@@ -190,15 +189,8 @@ impl NodePvm {
             }
         }
 
-        let refs = proof_state.struct_ref::<FnManagerIdent>();
-        let merkle_proof = NodePvmLayout::to_merkle_tree(refs)
-            .ok()?
-            .to_merkle_proof()
-            .ok()?;
-
-        let refs = proof_state.struct_ref::<FnManagerIdent>();
-        let final_hash = NodePvmLayout::state_hash(refs).ok()?;
-        Some(Proof::new(merkle_proof, final_hash))
+        let proof = proof_state.to_proof().ok()?;
+        Some(proof)
     }
 }
 
diff --git a/src/riscv/lib/src/stepper/pvm.rs b/src/riscv/lib/src/stepper/pvm.rs
index 0160994fc449..a0e1cd1efeb1 100644
--- a/src/riscv/lib/src/stepper/pvm.rs
+++ b/src/riscv/lib/src/stepper/pvm.rs
@@ -32,7 +32,6 @@ use crate::pvm::PvmLayout;
 use crate::pvm::PvmStatus;
 use crate::range_utils::bound_saturating_sub;
 use crate::state_backend::AllocatedOf;
-use crate::state_backend::CommitmentLayout;
 use crate::state_backend::FnManagerIdent;
 use crate::state_backend::ManagerBase;
 use crate::state_backend::ManagerReadWrite;
@@ -137,18 +136,10 @@ impl<MC: MemoryConfig, CL: CacheLayouts> PvmStepper<'_, MC, CL, Owned> {
         // Step using the proof mode stepper in order to obtain the proof
         let mut proof_stepper = self.start_proof_mode();
 
-        proof_stepper.try_step().then(|| {
-            let refs = proof_stepper.struct_ref();
-            let merkle_proof = PvmLayout::<MC, CL>::to_merkle_tree(refs)
-                .expect("Obtaining the Merkle tree of a proof-gen state should not fail")
-                .to_merkle_proof()
-                .expect("Converting a Merkle tree to a compressed Merkle proof should not fail");
-
-            let refs = proof_stepper.struct_ref();
-            let final_hash = PvmLayout::<MC, CL>::state_hash(refs)
-                .expect("Obtaining the final hash of the proof-gen state should not fail");
-            Proof::new(merkle_proof, final_hash)
-        })
+        proof_stepper.try_step().then_some(())?;
+
+        let proof = proof_stepper.pvm.to_proof().ok()?;
+        Some(proof)
     }
 }
 
-- 
GitLab