comments

crates/sdk/src/keygen/dummy.rs

@@ -55,7 +55,9 @@ pub(super) fn compute_root_proof_heights(
         public_values: vec![F::ZERO; num_user_public_values],
     };
     let vm = SingleSegmentVmExecutor::new(root_vm_config);
-    let res = vm.execute(root_exe, root_input.write(), true).unwrap();
+    let res = vm
+        .execute_and_compute_heights(root_exe, root_input.write())
+        .unwrap();
     let air_heights: Vec<_> = res
         .air_heights
         .into_iter()

crates/sdk/src/prover/root.rs

@@ -33,10 +33,9 @@ impl RootVerifierLocalProver {
     pub fn execute_for_air_heights(&self, input: RootVmVerifierInput<SC>) -> Vec<usize> {
         let result = self
             .executor_for_heights
-            .execute(
+            .execute_and_compute_heights(
                 self.root_verifier_pk.root_committed_exe.exe.clone(),
                 input.write(),
-                true,
             )
             .unwrap();
         result.air_heights

crates/sdk/tests/integration_test.rs

@@ -50,10 +50,9 @@ where
     VC::Executor: Chip<SC>,
     VC::Periphery: Chip<SC>,
 {
-    let exe_result = leaf_vm.execute(
+    let exe_result = leaf_vm.execute_and_compute_heights(
         leaf_committed_exe.exe.clone(),
         verifier_input.write_to_stream(),
-        false,
     )?;
     let runtime_pvs: Vec<_> = exe_result
         .public_values

crates/vm/src/arch/extensions.rs

@@ -1027,7 +1027,7 @@ impl<F: PrimeField32, E, P> VmChipComplex<F, E, P> {
     }
 
     #[cfg(feature = "bench-metrics")]
-    fn finalize_metrics(&mut self, metrics: &mut VmMetrics)
+    fn finalize_metrics(&self, metrics: &mut VmMetrics)
     where
         E: ChipUsageGetter,
         P: ChipUsageGetter,

crates/vm/src/arch/vm.rs

@@ -292,17 +292,17 @@ where
         self.overridden_heights = Some(overridden_heights);
     }
 
-    /// Executes a program and returns the public values. None means the public value is not set.
-    pub fn execute(
+    /// Executes a program, compute the trace heights, and returns the public values.
+    pub fn execute_and_compute_heights(
         &self,
         exe: impl Into<VmExe<F>>,
         input: impl Into<Streams<F>>,
-        compute_heights: bool,
     ) -> Result<SingleSegmentVmExecutionResult<F>, ExecutionError> {
-        let mut segment = self.execute_impl(exe.into(), input)?;
-        if compute_heights {
+        let segment = {
+            let mut segment = self.execute_impl(exe.into(), input)?;
             segment.chip_complex.finalize_memory();
-        }
+            segment
+        };
         let air_heights = segment.chip_complex.current_trace_heights();
         let internal_heights = segment.chip_complex.get_internal_trace_heights();
         let public_values = if let Some(pv_chip) = segment.chip_complex.public_values_chip() {

crates/vm/src/system/memory/controller/mod.rs

@@ -682,6 +682,12 @@ impl<F: PrimeField32> MemoryController<F> {
         ret
     }
 
+    /// Returns a reference to the offline memory.
+    ///
+    /// Until `finalize` is called, the `OfflineMemory` does not contain useful state, and should
+    /// therefore not be used by any chip during execution. However, to obtain a reference to the
+    /// offline memory that will be useful in trace generation, a chip can call `offline_memory()`
+    /// and store the returned reference for later use.
     pub fn offline_memory(&self) -> Arc<Mutex<OfflineMemory<F>>> {
         self.offline_memory.clone()
     }

crates/vm/tests/integration_test.rs

@@ -163,7 +163,7 @@ fn test_vm_override_executor_height() {
 
     let executor = SingleSegmentVmExecutor::new(vm_config.clone());
     let res = executor
-        .execute(committed_exe.exe.clone(), vec![], true)
+        .execute_and_compute_heights(committed_exe.exe.clone(), vec![])
         .unwrap();
     // Memory trace heights are not computed during execution.
     assert_eq!(
@@ -305,7 +305,9 @@ fn test_vm_public_values() {
             vm.engine.config.pcs(),
         ));
         let single_vm = SingleSegmentVmExecutor::new(config);
-        let exe_result = single_vm.execute(program, vec![], false).unwrap();
+        let exe_result = single_vm
+            .execute_and_compute_heights(program, vec![])
+            .unwrap();
         assert_eq!(
             exe_result.public_values,
             [

extensions/native/compiler/tests/public_values.rs

@@ -33,7 +33,9 @@ fn test_compiler_public_values() {
         Native,
     ));
 
-    let exe_result = executor.execute(program, vec![], false).unwrap();
+    let exe_result = executor
+        .execute_and_compute_heights(program, vec![])
+        .unwrap();
     assert_eq!(
         exe_result
             .public_values