From 657ec4480f9b6d9c839eaaca4fe6011b00b46768 Mon Sep 17 00:00:00 2001
From: Victor Dumitrescu <v.dumitrescu@outlook.com>
Date: Fri, 13 Dec 2024 10:29:29 +0100
Subject: [PATCH 1/4] RISC-V: Adapt ProofGen managers trait bounds

---
 src/riscv/lib/src/state_backend/proof_backend.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/riscv/lib/src/state_backend/proof_backend.rs b/src/riscv/lib/src/state_backend/proof_backend.rs
index 725adb40995f..36213b8e7a5f 100644
--- a/src/riscv/lib/src/state_backend/proof_backend.rs
+++ b/src/riscv/lib/src/state_backend/proof_backend.rs
@@ -148,7 +148,7 @@ impl<M: ManagerRead> ManagerRead for ProofGen<M> {
 
 /// Implementation of [`ManagerWrite`] which wraps another manager and
 /// records written locations but does not write to the wrapped region directly.
-impl<M: ManagerWrite> ManagerWrite for ProofGen<M> {
+impl<M: ManagerBase> ManagerWrite for ProofGen<M> {
     fn region_write<E, const LEN: usize>(
         region: &mut Self::Region<E, LEN>,
         index: usize,
@@ -211,7 +211,7 @@ impl<M: ManagerWrite> ManagerWrite for ProofGen<M> {
 
 /// Implementation of [`ManagerReadWrite`] which wraps another manager and
 /// additionally records read and written locations.
-impl<M: ManagerReadWrite> ManagerReadWrite for ProofGen<M> {
+impl<M: ManagerRead> ManagerReadWrite for ProofGen<M> {
     fn region_replace<E: Copy, const LEN: usize>(
         region: &mut Self::Region<E, LEN>,
         index: usize,
-- 
GitLab


From f363b7294a24ce59d11985b9fe1fe7bb6b72c812 Mon Sep 17 00:00:00 2001
From: Victor Dumitrescu <v.dumitrescu@outlook.com>
Date: Fri, 13 Dec 2024 10:57:27 +0100
Subject: [PATCH 2/4] RISC-V: Adapt Merkleise implementations

---
 .../lib/src/machine_state/main_memory.rs      | 10 ++++---
 .../lib/src/state_backend/proof_backend.rs    | 28 +++++++++++--------
 src/riscv/lib/src/state_backend/region.rs     | 21 ++++++++------
 3 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/src/riscv/lib/src/machine_state/main_memory.rs b/src/riscv/lib/src/machine_state/main_memory.rs
index d5256a6e99c3..fdf83d63f9d7 100644
--- a/src/riscv/lib/src/machine_state/main_memory.rs
+++ b/src/riscv/lib/src/machine_state/main_memory.rs
@@ -12,7 +12,7 @@ use crate::{
             merkle::{MerkleTree, Merkleisable},
             ProofGen,
         },
-        verify_backend, ManagerDeserialise, ManagerRead, ManagerSerialise,
+        verify_backend, ManagerDeserialise, ManagerRead, ManagerSerialise, Ref,
     },
     storage::binary,
 };
@@ -129,7 +129,7 @@ pub trait MainMemoryLayout: backend::ProofLayout + 'static {
     fn hash_data<M: ManagerRead>(data: &Self::Data<M>) -> Result<Hash, HashError>;
 
     fn to_merkle_tree<M: backend::ManagerRead>(
-        data: &Self::Data<ProofGen<M>>,
+        data: &Self::Data<Ref<'_, ProofGen<M>>>,
     ) -> Result<MerkleTree, HashError>;
 }
 
@@ -203,7 +203,7 @@ impl<const BYTES: usize> MainMemoryLayout for Sizes<BYTES> {
     }
 
     fn to_merkle_tree<M: backend::ManagerRead>(
-        data: &Self::Data<ProofGen<M>>,
+        data: &Self::Data<Ref<'_, ProofGen<M>>>,
     ) -> Result<MerkleTree, HashError> {
         data.to_merkle_tree()
     }
@@ -394,7 +394,9 @@ impl<L: MainMemoryLayout, M: ManagerRead> RootHashable for MainMemory<L, M> {
     }
 }
 
-impl<L: MainMemoryLayout, M: backend::ManagerRead> Merkleisable for MainMemory<L, ProofGen<M>> {
+impl<L: MainMemoryLayout, M: backend::ManagerRead> Merkleisable
+    for MainMemory<L, Ref<'_, ProofGen<M>>>
+{
     fn to_merkle_tree(&self) -> Result<MerkleTree, HashError> {
         L::to_merkle_tree(&self.data)
     }
diff --git a/src/riscv/lib/src/state_backend/proof_backend.rs b/src/riscv/lib/src/state_backend/proof_backend.rs
index 36213b8e7a5f..03d9e8bc620d 100644
--- a/src/riscv/lib/src/state_backend/proof_backend.rs
+++ b/src/riscv/lib/src/state_backend/proof_backend.rs
@@ -425,7 +425,7 @@ mod tests {
         owned_backend::Owned,
         proof_backend::merkle::Merkleisable,
         region::{DynCells, MERKLE_LEAF_SIZE},
-        Cells, EnrichedCell,
+        Cells, EnrichedCell, FnManagerIdent,
     };
     use proptest::{array, prop_assert_eq, proptest};
     use std::collections::VecDeque;
@@ -499,7 +499,10 @@ mod tests {
             let initial_root_hash = cells.hash().unwrap();
             cells.write(i, value_after);
             prop_assert_eq!(cells.hash().unwrap(), initial_root_hash);
-            let merkle_tree = cells.to_merkle_tree().unwrap();
+            let merkle_tree = cells.struct_ref::<FnManagerIdent>()
+                .to_merkle_tree()
+                .unwrap();
+
             match merkle_tree {
                 MerkleTree::Leaf(hash, access_info, _) => {
                     prop_assert_eq!(hash, initial_root_hash);
@@ -529,8 +532,7 @@ mod tests {
                     bytes_after: [u8; ELEM_SIZE],
                     write_address in &address_range)| {
             let cells = Box::new([byte_before; DYN_REGION_SIZE]);
-            let dyn_region: ProofDynRegion<DYN_REGION_SIZE, Owned> =
-                ProofDynRegion::bind(cells);
+            let dyn_region: ProofDynRegion<DYN_REGION_SIZE, Owned> = ProofDynRegion::bind(cells);
             let mut dyn_cells: DynCells<DYN_REGION_SIZE, ProofGen<Owned>> =
                 DynCells::bind(dyn_region);
 
@@ -545,8 +547,7 @@ mod tests {
             assert_eq!(value, value_after);
 
             let cells = Box::new([byte_before; DYN_REGION_SIZE]);
-            let dyn_region: ProofDynRegion<DYN_REGION_SIZE, Owned> =
-                ProofDynRegion::bind(cells);
+            let dyn_region: ProofDynRegion<DYN_REGION_SIZE, Owned> = ProofDynRegion::bind(cells);
             let mut dyn_cells: DynCells<DYN_REGION_SIZE, ProofGen<Owned>> =
                 DynCells::bind(dyn_region);
 
@@ -567,8 +568,7 @@ mod tests {
             assert_eq!(value, value_after);
 
             let cells = Box::new([byte_before; DYN_REGION_SIZE]);
-            let dyn_region: ProofDynRegion<DYN_REGION_SIZE, Owned> =
-                ProofDynRegion::bind(cells);
+            let dyn_region: ProofDynRegion<DYN_REGION_SIZE, Owned> = ProofDynRegion::bind(cells);
             let mut dyn_cells: DynCells<DYN_REGION_SIZE, ProofGen<Owned>> =
                 DynCells::bind(dyn_region);
 
@@ -607,8 +607,11 @@ mod tests {
 
             // Build the Merkle tree and check that it has the root hash of the
             // initial wrapped region.
-            let merkle_tree = dyn_cells.to_merkle_tree().unwrap();
-            let cells: DynCells<DYN_REGION_SIZE,Owned> = DynCells::bind(region);
+            let merkle_tree = dyn_cells
+                .struct_ref::<FnManagerIdent>()
+                .to_merkle_tree()
+                .unwrap();
+            let cells: DynCells<DYN_REGION_SIZE, Owned> = DynCells::bind(region);
             let initial_root_hash = cells.hash().unwrap();
             prop_assert_eq!(merkle_tree.root_hash(), initial_root_hash);
 
@@ -698,7 +701,10 @@ mod tests {
             let initial_root_hash = proof_cell.hash().unwrap();
             proof_cell.write(value_after);
             prop_assert_eq!(proof_cell.hash().unwrap(), initial_root_hash);
-            let merkle_tree = proof_cell.to_merkle_tree().unwrap();
+            let merkle_tree = proof_cell
+                .struct_ref::<FnManagerIdent>()
+                .to_merkle_tree()
+                .unwrap();
             match merkle_tree {
                 MerkleTree::Leaf(hash, access_info, _) => {
                     prop_assert_eq!(hash, initial_root_hash);
diff --git a/src/riscv/lib/src/state_backend/region.rs b/src/riscv/lib/src/state_backend/region.rs
index 570f49b436e6..77812c0d3c61 100644
--- a/src/riscv/lib/src/state_backend/region.rs
+++ b/src/riscv/lib/src/state_backend/region.rs
@@ -198,13 +198,15 @@ impl<A: PartialEq<B> + Copy, B: Copy, M: ManagerRead, N: ManagerRead> PartialEq<
     }
 }
 
-impl<E: serde::Serialize, M: ManagerSerialise> Merkleisable for Cell<E, ProofGen<M>> {
+impl<E: serde::Serialize, M: ManagerSerialise> Merkleisable for Cell<E, Ref<'_, ProofGen<M>>> {
     fn to_merkle_tree(&self) -> Result<MerkleTree, HashError> {
         self.region.to_merkle_tree()
     }
 }
 
-impl<E: serde::Serialize, M: ManagerSerialise> AccessInfoAggregatable for Cell<E, ProofGen<M>> {
+impl<E: serde::Serialize, M: ManagerSerialise> AccessInfoAggregatable
+    for Cell<E, Ref<'_, ProofGen<M>>>
+{
     fn aggregate_access_info(&self) -> AccessInfo {
         self.region.region.get_access_info()
     }
@@ -416,17 +418,18 @@ where
     }
 }
 
-impl<V: EnrichedValue, M: ManagerSerialise> Merkleisable for EnrichedCell<V, ProofGen<M>>
+impl<V: EnrichedValue, M: ManagerSerialise> Merkleisable for EnrichedCell<V, Ref<'_, ProofGen<M>>>
 where
     V::E: serde::Serialize,
 {
     fn to_merkle_tree(&self) -> Result<MerkleTree, HashError> {
-        let serialised = ProofEnrichedCell::serialise_inner_enriched_cell(&self.cell)?;
+        let serialised = ProofEnrichedCell::serialise_inner_enriched_cell(self.cell)?;
         MerkleTree::make_merkle_leaf(serialised, self.cell.get_access_info())
     }
 }
 
-impl<V: EnrichedValue, M: ManagerSerialise> AccessInfoAggregatable for EnrichedCell<V, ProofGen<M>>
+impl<V: EnrichedValue, M: ManagerSerialise> AccessInfoAggregatable
+    for EnrichedCell<V, Ref<'_, ProofGen<M>>>
 where
     V::E: serde::Serialize,
 {
@@ -476,7 +479,7 @@ impl<A: PartialEq<B> + Copy, B: Copy, const LEN: usize, M: ManagerRead, N: Manag
 }
 
 impl<E: serde::Serialize, const LEN: usize, M: ManagerSerialise> Merkleisable
-    for Cells<E, LEN, ProofGen<M>>
+    for Cells<E, LEN, Ref<'_, ProofGen<M>>>
 {
     fn to_merkle_tree(&self) -> Result<MerkleTree, HashError> {
         // RV-282: Break down into multiple leaves if the size of the `Cells`
@@ -487,7 +490,7 @@ impl<E: serde::Serialize, const LEN: usize, M: ManagerSerialise> Merkleisable
 }
 
 impl<E: serde::Serialize, const LEN: usize, M: ManagerSerialise> AccessInfoAggregatable
-    for Cells<E, LEN, ProofGen<M>>
+    for Cells<E, LEN, Ref<'_, ProofGen<M>>>
 {
     fn aggregate_access_info(&self) -> AccessInfo {
         self.region.get_access_info()
@@ -634,7 +637,7 @@ impl<const LEN: usize, M: ManagerRead> RootHashable for DynCells<LEN, M> {
     }
 }
 
-impl<const LEN: usize, M: ManagerRead> Merkleisable for DynCells<LEN, ProofGen<M>> {
+impl<const LEN: usize, M: ManagerRead> Merkleisable for DynCells<LEN, Ref<'_, ProofGen<M>>> {
     fn to_merkle_tree(&self) -> Result<MerkleTree, HashError> {
         let mut writer = MerkleWriter::new(
             MERKLE_LEAF_SIZE,
@@ -644,7 +647,7 @@ impl<const LEN: usize, M: ManagerRead> Merkleisable for DynCells<LEN, ProofGen<M
             LEN.div_ceil(MERKLE_ARITY),
         );
         let read = |address| -> [u8; MERKLE_LEAF_SIZE.get()] {
-            ProofDynRegion::inner_dyn_region_read(&self.region, address)
+            ProofDynRegion::inner_dyn_region_read(self.region, address)
         };
         chunks_to_writer::<LEN, _, _>(&mut writer, read)?;
         writer.finalise()
-- 
GitLab


From e14a2ee58aefa1d91126bf291b84f702246c43c9 Mon Sep 17 00:00:00 2001
From: Victor Dumitrescu <v.dumitrescu@outlook.com>
Date: Thu, 19 Dec 2024 08:57:27 +0100
Subject: [PATCH 3/4] RISC-V: Proof exposes initial state hash

---
 .../src/state_backend/proof_backend/merkle.rs |  4 +++
 .../src/state_backend/proof_backend/proof.rs  | 34 +++++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/riscv/lib/src/state_backend/proof_backend/merkle.rs b/src/riscv/lib/src/state_backend/proof_backend/merkle.rs
index 243b35635168..88abae1465b7 100644
--- a/src/riscv/lib/src/state_backend/proof_backend/merkle.rs
+++ b/src/riscv/lib/src/state_backend/proof_backend/merkle.rs
@@ -754,6 +754,10 @@ mod tests {
             assert_eq!(compressed_merkle_proof, proof);
 
             assert_eq!(merkle_tree.to_merkle_proof().unwrap(), proof);
+            assert_eq!(
+                compressed_merkle_proof.hash().unwrap(),
+                merkle_tree_root_hash
+            );
 
             Ok(())
         };
diff --git a/src/riscv/lib/src/state_backend/proof_backend/proof.rs b/src/riscv/lib/src/state_backend/proof_backend/proof.rs
index 876ad19fe080..46699db9d54e 100644
--- a/src/riscv/lib/src/state_backend/proof_backend/proof.rs
+++ b/src/riscv/lib/src/state_backend/proof_backend/proof.rs
@@ -1,4 +1,5 @@
 // SPDX-FileCopyrightText: 2024 TriliTech <contact@trili.tech>
+// SPDX-FileCopyrightText: 2024 Nomadic Labs <contact@nomadic-labs.com>
 //
 // SPDX-License-Identifier: MIT
 
@@ -13,8 +14,11 @@
 //!
 //! - Convert [`MerkleTree`] to [`MerkleProof`]
 
-use super::tree::{ModifyResult, Tree};
-use crate::{state_backend::hash::Hash, storage::DIGEST_SIZE};
+use super::tree::{impl_modify_map_collect, ModifyResult, Tree};
+use crate::{
+    state_backend::hash::{Hash, RootHashable},
+    storage::{HashError, DIGEST_SIZE},
+};
 use itertools::Itertools;
 
 /// Structure of a proof transitioning from state A to state B.
@@ -45,6 +49,13 @@ impl Proof {
         &self.partial_tree
     }
 
+    /// Get the initial state hash of the proof.
+    pub fn initial_state_hash(&self) -> Hash {
+        self.partial_tree
+            .hash()
+            .expect("Computing the root hash of the Merkle proof should not fail")
+    }
+
     /// Get the final state hash of the proof.
     pub fn final_state_hash(&self) -> &Hash {
         &self.final_state_hash
@@ -99,6 +110,25 @@ impl MerkleProof {
     }
 }
 
+impl RootHashable for MerkleProof {
+    fn hash(&self) -> Result<Hash, HashError> {
+        impl_modify_map_collect(
+            self,
+            |subtree| {
+                Ok(match subtree {
+                    Tree::Node(vec) => ModifyResult::NodeContinue((), vec.iter().collect()),
+                    Tree::Leaf(data) => ModifyResult::LeafStop(data),
+                })
+            },
+            |leaf| match leaf {
+                MerkleProofLeaf::Blind(hash) => Ok(*hash),
+                MerkleProofLeaf::Read(data) => Hash::blake2b_hash_bytes(data.as_slice()),
+            },
+            |(), leaves| leaves.hash(),
+        )
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub enum Tag {
     Node,
-- 
GitLab


From 9aeb5e5c1d23ac3265e3da1b76c0aaf3b9f55c70 Mon Sep 17 00:00:00 2001
From: Victor Dumitrescu <v.dumitrescu@outlook.com>
Date: Wed, 18 Dec 2024 09:30:48 +0100
Subject: [PATCH 4/4] RISC-V: Stepper can produce proof

---
 src/riscv/lib/src/pvm/common.rs         | 35 +++++++++++++-
 src/riscv/lib/src/stepper/pvm.rs        | 62 ++++++++++++++++++++++---
 src/riscv/lib/tests/common/mod.rs       |  6 +--
 src/riscv/lib/tests/test_determinism.rs |  2 +-
 src/riscv/lib/tests/test_proofs.rs      | 36 +++++++++++---
 5 files changed, 122 insertions(+), 19 deletions(-)

diff --git a/src/riscv/lib/src/pvm/common.rs b/src/riscv/lib/src/pvm/common.rs
index f76a036db2b1..d53490f0a698 100644
--- a/src/riscv/lib/src/pvm/common.rs
+++ b/src/riscv/lib/src/pvm/common.rs
@@ -6,7 +6,11 @@ use crate::{
     default::ConstDefault,
     machine_state::{self, main_memory},
     pvm::sbi,
-    state_backend::{self, Atom, Cell},
+    state_backend::{
+        self,
+        proof_backend::{ProofDynRegion, ProofEnrichedCell, ProofGen, ProofRegion},
+        Atom, Cell,
+    },
     traps::EnvironException,
 };
 use std::{
@@ -134,6 +138,35 @@ impl<
         )
     }
 
+    /// Generate a proof-generating version of this PVM.
+    pub fn start_proof(&self) -> Pvm<ML, CL, ProofGen<state_backend::Ref<'_, M>>> {
+        enum ProofWrapper {}
+
+        impl<M: state_backend::ManagerBase> state_backend::FnManager<M> for ProofWrapper {
+            type Output = ProofGen<M>;
+
+            fn map_region<E: 'static, const LEN: usize>(
+                input: <M as state_backend::ManagerBase>::Region<E, LEN>,
+            ) -> <ProofGen<M> as state_backend::ManagerBase>::Region<E, LEN> {
+                ProofRegion::bind(input)
+            }
+
+            fn map_dyn_region<const LEN: usize>(
+                input: <M as state_backend::ManagerBase>::DynRegion<LEN>,
+            ) -> <ProofGen<M> as state_backend::ManagerBase>::DynRegion<LEN> {
+                ProofDynRegion::bind(input)
+            }
+
+            fn map_enriched_cell<V: state_backend::EnrichedValue>(
+                input: <M as state_backend::ManagerBase>::EnrichedCell<V>,
+            ) -> <ProofGen<M> as state_backend::ManagerBase>::EnrichedCell<V> {
+                ProofEnrichedCell::bind(input)
+            }
+        }
+
+        Pvm::bind(self.struct_ref::<ProofWrapper>())
+    }
+
     /// Reset the PVM state.
     pub fn reset(&mut self)
     where
diff --git a/src/riscv/lib/src/stepper/pvm.rs b/src/riscv/lib/src/stepper/pvm.rs
index 59f511ea95b6..8c2bf0e9106c 100644
--- a/src/riscv/lib/src/stepper/pvm.rs
+++ b/src/riscv/lib/src/stepper/pvm.rs
@@ -1,8 +1,10 @@
 // SPDX-FileCopyrightText: 2024 TriliTech <contact@trili.tech>
+// SPDX-FileCopyrightText: 2024 Nomadic Labs <contact@nomadic-labs.com>
 //
 // SPDX-License-Identifier: MIT
 
 use super::{Stepper, StepperStatus};
+use crate::state_backend::{ManagerBase, ManagerReadWrite};
 use crate::{
     kernel_loader,
     machine_state::{
@@ -16,7 +18,7 @@ use crate::{
     state_backend::{
         hash::{Hash, RootHashable},
         owned_backend::Owned,
-        proof_backend::proof::Proof,
+        proof_backend::{merkle::Merkleisable, proof::Proof, ProofGen},
         verify_backend::Verifier,
         AllocatedOf, FnManagerIdent, ProofLayout, ProofTree, Ref,
     },
@@ -37,8 +39,13 @@ pub enum PvmStepperError {
 }
 
 /// Wrapper over a PVM that lets you step through it
-pub struct PvmStepper<'hooks, ML: MainMemoryLayout = M1G, CL: CacheLayouts = DefaultCacheLayouts> {
-    pvm: Pvm<ML, CL, Owned>,
+pub struct PvmStepper<
+    'hooks,
+    ML: MainMemoryLayout = M1G,
+    CL: CacheLayouts = DefaultCacheLayouts,
+    M: ManagerBase = Owned,
+> {
+    pvm: Pvm<ML, CL, M>,
     hooks: PvmHooks<'hooks>,
     inbox: Inbox,
     rollup_address: [u8; 20],
@@ -135,10 +142,16 @@ impl<'hooks, ML: MainMemoryLayout, CL: CacheLayouts> PvmStepper<'hooks, ML, CL>
         RootHashable::hash(&refs).unwrap()
     }
 
-    /// Produce the Merkle proof for evaluating one step on the given PVM state.
-    pub fn produce_proof(&mut self) -> Proof {
-        // TODO RV-338 PVM stepper can produce a proof
-        todo!()
+    /// Create a new stepper in which the existing PVM is managed by
+    /// the proof-generating backend.
+    pub fn start_proof_mode<'a>(&'a self) -> PvmStepper<'hooks, ML, CL, ProofGen<Ref<'a, Owned>>> {
+        PvmStepper::<'hooks, ML, CL, ProofGen<Ref<'a, Owned>>> {
+            pvm: self.pvm.start_proof(),
+            hooks: PvmHooks::none(),
+            inbox: self.inbox.clone(),
+            rollup_address: self.rollup_address,
+            origination_level: self.origination_level,
+        }
     }
 
     /// Verify a Merkle proof. The [`PvmStepper`] is used for inbox information.
@@ -191,6 +204,41 @@ impl<'hooks, ML: MainMemoryLayout, CL: CacheLayouts> PvmStepper<'hooks, ML, CL>
     }
 }
 
+impl<'hooks, ML: MainMemoryLayout, CL: CacheLayouts, M: ManagerReadWrite>
+    PvmStepper<'hooks, ML, CL, M>
+{
+    /// Perform one evaluation step.
+    pub fn eval_one(&mut self) {
+        self.pvm.eval_one(&mut self.hooks)
+    }
+}
+
+impl<'hooks, ML: MainMemoryLayout, CL: CacheLayouts> PvmStepper<'hooks, ML, CL>
+where
+    for<'a, 'b> AllocatedOf<PvmLayout<ML, CL>, Ref<'a, ProofGen<Ref<'b, Owned>>>>: Merkleisable,
+    for<'a> AllocatedOf<PvmLayout<ML, CL>, Ref<'a, Owned>>: RootHashable,
+{
+    /// Produce the Merkle proof for evaluating one step on the given PVM state.
+    /// The given stepper takes one step.
+    pub fn produce_proof(&mut self) -> Option<Proof> {
+        // Step using the proof mode stepper in order to obtain the proof
+        let mut proof_stepper = self.start_proof_mode();
+        proof_stepper.eval_one();
+        let merkle_proof = proof_stepper
+            .pvm
+            .struct_ref::<crate::state_backend::FnManagerIdent>()
+            .to_merkle_tree()
+            .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");
+
+        // TODO RV-373 : Proof-generating backend should also compute final state hash
+        // Currently, the proof and the initial state hash are valid,
+        // but the final state hash is not.
+        Some(Proof::new(merkle_proof, self.hash()))
+    }
+}
+
 impl<'hooks, ML: MainMemoryLayout, CL: CacheLayouts> Stepper for PvmStepper<'hooks, ML, CL> {
     type MainMemoryLayout = ML;
 
diff --git a/src/riscv/lib/tests/common/mod.rs b/src/riscv/lib/tests/common/mod.rs
index be2912915c8b..30a4a8a85368 100644
--- a/src/riscv/lib/tests/common/mod.rs
+++ b/src/riscv/lib/tests/common/mod.rs
@@ -40,15 +40,15 @@ pub fn make_stepper_factory() -> impl Fn() -> PvmStepper<'static, M100M, Default
     }
 }
 
-pub fn dissect_steps(mut max_steps: usize) -> Vec<usize> {
+pub fn dissect_steps(mut max_steps: usize, min_step_size: usize) -> Vec<usize> {
     let mut rng = rand::thread_rng();
     let mut steps: Vec<usize> = std::iter::from_fn(|| {
         if max_steps == 0 {
             return None;
         }
 
-        let steps = max_steps.div_euclid(2).max(1);
-        let steps = rng.gen_range(0..=steps);
+        let steps = max_steps.div_euclid(2).max(min_step_size + 1);
+        let steps = rng.gen_range(min_step_size..=steps);
 
         max_steps = max_steps.saturating_sub(steps);
 
diff --git a/src/riscv/lib/tests/test_determinism.rs b/src/riscv/lib/tests/test_determinism.rs
index b9bf80d63a9b..d51e6b5e70cb 100644
--- a/src/riscv/lib/tests/test_determinism.rs
+++ b/src/riscv/lib/tests/test_determinism.rs
@@ -37,7 +37,7 @@ fn test_jstz_determinism() {
     let base_refs = base_stepper.struct_ref();
 
     // Create multiple series of bisections that we will evaluate.
-    let ladder = dissect_steps(steps);
+    let ladder = dissect_steps(steps, 0);
     run_steps_ladder(&make_stepper, &ladder, &base_refs, base_hash);
 }
 
diff --git a/src/riscv/lib/tests/test_proofs.rs b/src/riscv/lib/tests/test_proofs.rs
index b0c9e7e8e324..ee21b993efbb 100644
--- a/src/riscv/lib/tests/test_proofs.rs
+++ b/src/riscv/lib/tests/test_proofs.rs
@@ -25,7 +25,9 @@ fn test_jstz_proofs() {
     let steps = base_result.steps();
     let base_hash = base_stepper.hash();
 
-    let ladder = dissect_steps(steps);
+    // For each step `s`, the stepper will initially step `s-1` steps, then
+    // produce a proof of the `s` step. The minimum step size thus needs to be 1.
+    let ladder = dissect_steps(steps, 1);
     run_steps_ladder(&make_stepper, &ladder, base_hash);
 }
 
@@ -38,19 +40,39 @@ where
 
     let mut steps_done = 0;
     for &steps in ladder {
+        // Run one step short of `steps`, then produce a proof of the following step.
+        let steps = steps.checked_sub(1).expect("minimum step size is 1");
         eprintln!("> Running {} steps ...", steps);
-        let StepperStatus::Running { .. } = stepper.step_max(Bound::Included(steps)) else {
-            panic!("Unexpected stepper result")
-        };
+        let result = stepper.step_max(Bound::Included(steps));
         steps_done += steps;
 
+        if steps_done != expected_steps {
+            assert!(matches!(result, StepperStatus::Running { .. }));
+
+            eprintln!("> Producing proof");
+            let proof = stepper.produce_proof().unwrap();
+
+            assert_eq!(proof.initial_state_hash(), stepper.hash());
+
+            // Run one final step, which is the step proven by `proof`, and check that its
+            // state hash matches the final state hash of `proof`.
+            eprintln!("> Running 1 step ...");
+            stepper.eval_one();
+            steps_done += 1;
+            // TODO RV-373 : Proof-generating backend should also compute final state hash
+            assert_eq!(proof.final_state_hash(), &stepper.hash());
+
+            assert!(stepper.verify_proof(proof))
+        } else {
+            // Can't generate a proof on the next step if execution has ended
+            assert!(matches!(result, StepperStatus::Exited { .. }));
+            assert!(ladder.is_empty())
+        };
+
         eprintln!(
             "> Done {:.2}%",
             (steps_done as f64 / expected_steps as f64) * 100.0
         );
-
-        let proof = stepper.produce_proof();
-        assert!(stepper.verify_proof(proof))
     }
 
     assert_eq!(stepper.hash(), expected_hash);
-- 
GitLab