feat(prt-client): improve snapshots

prt/client-lua/computation/commitment.lua

@@ -6,7 +6,6 @@ local consts = require "computation.constants"
 
 local ulte = arithmetic.ulte
 
-local save_snapshot = true
 local handle_rollups = false
 
 
@@ -48,12 +47,7 @@ local function run_uarch_span(machine)
     return builder:build(), machine_state
 end
 
-local function build_small_machine_commitment(base_cycle, log2_stride_count, machine, initial_state, snapshot_dir)
-    if save_snapshot then
-        -- taking snapshot for leafs to save time in next level
-        machine:take_snapshot(snapshot_dir, base_cycle, handle_rollups)
-    end
-
+local function build_small_machine_commitment(log2_stride_count, machine, initial_state, snapshot_dir)
     local builder = MerkleBuilder:new()
     local instruction_count = arithmetic.max_uint(log2_stride_count - consts.log2_uarch_span)
     local instruction = 0
@@ -81,13 +75,7 @@ local function build_small_machine_commitment(base_cycle, log2_stride_count, mac
     return initial_state, builder:build(initial_state)
 end
 
-local function build_big_machine_commitment(base_cycle, log2_stride, log2_stride_count, machine, initial_state,
-                                            snapshot_dir)
-    if save_snapshot then
-        -- taking snapshot for leafs to save time in next level
-        machine:take_snapshot(snapshot_dir, base_cycle, handle_rollups)
-    end
-
+local function build_big_machine_commitment(base_cycle, log2_stride, log2_stride_count, machine, initial_state)
     local builder = MerkleBuilder:new()
     local instruction_count = arithmetic.max_uint(log2_stride_count)
     local instruction = 0
@@ -125,11 +113,12 @@ local function build_commitment(base_cycle, log2_stride, log2_stride_count, mach
         -- the base_cycle may be the cycle to receive input,
         -- we need to take the initial state before feeding input to the machine
         handle_rollups = true
-        initial_state = machine:run_with_inputs(base_cycle, inputs).root_hash
+        initial_state = machine:run_with_inputs(base_cycle, inputs, snapshot_dir).root_hash
     else
         -- treat it as compute
         handle_rollups = false
-        initial_state = machine:run(base_cycle).root_hash
+        initial_state = machine:run(base_cycle).root_hash -- taking snapshot for leafs to save time in next level
+        machine:take_snapshot(snapshot_dir, base_cycle, handle_rollups)
     end
 
     if log2_stride >= consts.log2_uarch_span then
@@ -141,7 +130,7 @@ local function build_commitment(base_cycle, log2_stride, log2_stride_count, mach
             snapshot_dir)
     else
         assert(log2_stride == 0)
-        return build_small_machine_commitment(base_cycle, log2_stride_count, machine, initial_state, snapshot_dir)
+        return build_small_machine_commitment(log2_stride_count, machine, initial_state, snapshot_dir)
     end
 end
 

prt/client-lua/computation/machine.lua

@@ -117,9 +117,9 @@ end
 
 function Machine:take_snapshot(snapshot_dir, cycle, handle_rollups)
     local input_mask = arithmetic.max_uint(consts.log2_emulator_span)
-    if handle_rollups and cycle & input_mask == 0 and not self.yielded then
+    if handle_rollups and cycle & input_mask == 0 then
         -- dont snapshot a machine state that's freshly fed with input without advance
-        return
+        assert(not self.yielded, "don't snapshot a machine state that's freshly fed with input without advance")
     end
 
     if helper.exists(snapshot_dir) then
@@ -178,7 +178,7 @@ function Machine:run_uarch(ucycle)
     self.ucycle = ucycle
 end
 
-function Machine:run_with_inputs(cycle, inputs)
+function Machine:run_with_inputs(cycle, inputs, snapshot_dir)
     local input_mask = arithmetic.max_uint(consts.log2_emulator_span)
     local current_input_index = self.cycle >> consts.log2_emulator_span
 
@@ -194,6 +194,9 @@ function Machine:run_with_inputs(cycle, inputs)
 
     while next_input_cycle <= cycle do
         machine_state_without_input = self:run(next_input_cycle)
+        if next_input_cycle == cycle then
+            self:take_snapshot(snapshot_dir, next_input_cycle, true)
+        end
         local input = inputs[next_input_index + 1]
         if input then
             local h = assert(input:match("0x(%x+)"), input)
@@ -209,6 +212,7 @@ function Machine:run_with_inputs(cycle, inputs)
 
     if cycle > self.cycle then
         machine_state_without_input = self:run(cycle)
+        self:take_snapshot(snapshot_dir, next_input_cycle, true)
     end
 
     return machine_state_without_input

prt/client-rs/src/machine/commitment.rs

@@ -90,8 +90,6 @@ pub fn build_big_machine_commitment(
     initial_state: Digest,
     db: &ComputeStateAccess,
 ) -> Result<MachineCommitment> {
-    snapshot_base_cycle(machine, base_cycle, db)?;
-
     let mut builder = MerkleBuilder::default();
     let mut leafs = Vec::new();
     let instruction_count = arithmetic::max_uint(log2_stride_count);
@@ -159,8 +157,6 @@ pub fn build_small_machine_commitment(
     initial_state: Digest,
     db: &ComputeStateAccess,
 ) -> Result<MachineCommitment> {
-    snapshot_base_cycle(machine, base_cycle, db)?;
-
     let mut builder = MerkleBuilder::default();
     let mut leafs = Vec::new();
     let mut uarch_span_and_leafs = Vec::new();
@@ -208,22 +204,6 @@ pub fn build_small_machine_commitment(
     })
 }
 
-fn snapshot_base_cycle(
-    machine: &mut MachineInstance,
-    base_cycle: u64,
-    db: &ComputeStateAccess,
-) -> Result<()> {
-    let mask = arithmetic::max_uint(constants::LOG2_EMULATOR_SPAN);
-    if db.handle_rollups && base_cycle & mask == 0 && !machine.machine_state()?.yielded {
-        // don't snapshot a machine state that's freshly fed with input without advance
-        return Ok(());
-    }
-
-    let snapshot_path = db.work_path.join(format!("{}", base_cycle));
-    machine.snapshot(&snapshot_path)?;
-    Ok(())
-}
-
 fn run_uarch_span(
     machine: &mut MachineInstance,
 ) -> Result<(Arc<MerkleTree>, MachineState, Vec<Leaf>)> {

prt/client-rs/src/machine/commitment_builder.rs

@@ -48,12 +48,12 @@ impl CachingMachineCommitmentBuilder {
         let initial_state = {
             if db.handle_rollups {
                 // treat it as rollups
-                machine
-                    .run_with_inputs(base_cycle, &db.inputs()?)?
-                    .root_hash
+                machine.run_with_inputs(base_cycle, &db)?.root_hash
             } else {
                 // treat it as compute
-                machine.run(base_cycle)?.root_hash
+                let root_hash = machine.run(base_cycle)?.root_hash;
+                machine.take_snapshot(base_cycle, &db);
+                root_hash
             }
         };
         trace!("initial state for commitment: {}", initial_state);

prt/client-rs/src/machine/instance.rs

@@ -1,4 +1,4 @@
-use crate::db::compute_state_access::Input;
+use crate::db::compute_state_access::{ComputeStateAccess, Input};
 use crate::machine::constants;
 use cartesi_dave_arithmetic as arithmetic;
 use cartesi_dave_merkle::Digest;
@@ -65,6 +65,22 @@ impl MachineInstance {
             ucycle: 0,
         })
     }
+    pub fn take_snapshot(&mut self, base_cycle: u64, db: &ComputeStateAccess) -> Result<()> {
+        let mask = arithmetic::max_uint(constants::LOG2_EMULATOR_SPAN);
+        if db.handle_rollups && base_cycle & mask == 0 {
+            // don't snapshot a machine state that's freshly fed with input without advance
+            assert!(
+                self.machine_state()?.yielded,
+                "don't snapshot a machine state that's freshly fed with input without advance",
+            );
+        }
+
+        let snapshot_path = db.work_path.join(format!("{}", base_cycle));
+        if !snapshot_path.exists() {
+            self.machine.store(&snapshot_path)?;
+        }
+        Ok(())
+    }
 
     // load inner machine with snapshot, update cycle, keep everything else the same
     pub fn load_snapshot(&mut self, snapshot_path: &Path, snapshot_cycle: u64) -> Result<()> {
@@ -86,13 +102,6 @@ impl MachineInstance {
         Ok(())
     }
 
-    pub fn snapshot(&self, snapshot_path: &Path) -> Result<()> {
-        if !snapshot_path.exists() {
-            self.machine.store(snapshot_path)?;
-        }
-        Ok(())
-    }
-
     pub fn root_hash(&self) -> Digest {
         self.root_hash
     }
@@ -101,8 +110,7 @@ impl MachineInstance {
         &mut self,
         cycle: u64,
         ucycle: u64,
-        inputs: Vec<Vec<u8>>,
-        handle_rollups: bool,
+        db: &ComputeStateAccess,
     ) -> Result<MachineProof> {
         let log_type = AccessLogType {
             annotations: true,
@@ -111,18 +119,19 @@ impl MachineInstance {
         };
 
         let mut logs = Vec::new();
-        if handle_rollups {
+        if db.handle_rollups {
             // treat it as rollups
             // the cycle may be the cycle to receive input,
             // we need to include the process of feeding input to the machine in the log
             if cycle == 0 {
                 self.run(cycle)?;
             } else {
-                self.run_with_inputs(cycle - 1, &inputs)?;
+                self.run_with_inputs(cycle - 1, db)?;
                 self.run(cycle)?;
             }
 
             let mask = arithmetic::max_uint(constants::LOG2_EMULATOR_SPAN);
+            let inputs = &db.inputs()?;
             let input = inputs.get((cycle >> constants::LOG2_EMULATOR_SPAN) as usize);
             if cycle & mask == 0 && input.is_some() {
                 // need to process input
@@ -225,13 +234,14 @@ impl MachineInstance {
     // One exception is that if `cycle` is supposed to receive an input, in this case
     // the machine state would be `without` input included in the machine,
     // this is useful when we need the initial state to compute the commitments
-    pub fn run_with_inputs(&mut self, cycle: u64, inputs: &Vec<Vec<u8>>) -> Result<MachineState> {
+    pub fn run_with_inputs(&mut self, cycle: u64, db: &ComputeStateAccess) -> Result<MachineState> {
         trace!(
             "run_with_inputs self cycle: {}, target cycle: {}",
             self.cycle,
             cycle
         );
 
+        let inputs = &db.inputs()?;
         let mut machine_state_without_input = self.machine_state()?;
         let input_mask = arithmetic::max_uint(constants::LOG2_EMULATOR_SPAN);
         let current_input_index = self.cycle >> constants::LOG2_EMULATOR_SPAN;
@@ -249,15 +259,21 @@ impl MachineInstance {
             trace!("next input index: {}", next_input_index);
             trace!("run to next input cycle: {}", next_input_cycle);
             machine_state_without_input = self.run(next_input_cycle)?;
+            if next_input_cycle == cycle {
+                self.take_snapshot(next_input_cycle, &db)?;
+            }
+
             let input = inputs.get(next_input_index as usize);
             if let Some(data) = input {
                 trace!(
                     "before input, machine state: {}",
                     self.machine_state()?.root_hash
                 );
                 trace!("input: 0x{}", data.encode_hex());
+
                 self.machine
                     .send_cmio_response(htif::fromhost::ADVANCE_STATE, data)?;
+
                 trace!(
                     "after input, machine state: {}",
                     self.machine_state()?.root_hash
@@ -269,6 +285,7 @@ impl MachineInstance {
         }
         if cycle > self.cycle {
             machine_state_without_input = self.run(cycle)?;
+            self.take_snapshot(next_input_cycle, &db)?;
         }
         Ok(machine_state_without_input)
     }

prt/client-rs/src/strategy/player.rs

@@ -363,8 +363,7 @@ impl Player {
                 if let Some(snapshot) = self.db.closest_snapshot(cycle)? {
                     machine.load_snapshot(&snapshot.1, snapshot.0)?;
                 };
-                let inputs = self.db.inputs()?;
-                machine.get_logs(cycle, ucycle, inputs, self.db.handle_rollups)?
+                machine.get_logs(cycle, ucycle, &self.db)?
             };
 
             info!(