calc root heights

crates/sdk/src/keygen/dummy.rs

@@ -177,12 +177,14 @@ where
     } else {
         // We first execute once to get the trace heights from dummy_exe, then pad to powers of 2 (forcing trace height 0 to 1)
         let executor = VmExecutor::new(app_vm_pk.vm_config.clone());
-        let results = executor
+        let mut results = executor
             .execute_segments(dummy_exe.exe.clone(), vec![])
             .unwrap();
         // ASSUMPTION: the dummy exe has only 1 segment
         assert_eq!(results.len(), 1, "dummy exe should have only 1 segment");
-        let mut internal_heights = results[0].chip_complex.get_internal_trace_heights();
+        let mut result = results.pop().unwrap();
+        result.chip_complex.finalize_memory();
+        let mut internal_heights = result.chip_complex.get_internal_trace_heights();
         internal_heights.round_to_next_power_of_two();
         internal_heights
     };

crates/vm/src/arch/extensions.rs

@@ -714,6 +714,22 @@ impl<F: PrimeField32, E, P> VmChipComplex<F, E, P> {
         chip.as_any_kind_mut().downcast_mut()
     }
 
+    pub fn finalize_memory(&mut self)
+    where
+        P: AnyEnum,
+    {
+        let chip = self
+            .inventory
+            .periphery
+            .get_mut(VmChipComplex::<F, E, P>::POSEIDON2_PERIPHERY_IDX)
+            .expect("Poseidon2 chip required for persistent memory");
+        let hasher: &mut Poseidon2PeripheryChip<F> = chip
+            .as_any_kind_mut()
+            .downcast_mut()
+            .expect("Poseidon2 chip required for persistent memory");
+        self.base.memory_controller.finalize(Some(hasher));
+    }
+
     pub(crate) fn set_program(&mut self, program: Program<F>) {
         self.base.program_chip.set_program(program);
     }

crates/vm/src/arch/segment.rs

@@ -189,17 +189,13 @@ impl<F: PrimeField32, VC: VmConfig<F>> ExecutionSegment<F, VC> {
                 break;
             }
         }
-        // Finalize memory.
-        {
-            // Need some partial borrows, so code is ugly:
-            self.final_memory = Some(
-                self.chip_complex
-                    .base
-                    .memory_controller
-                    .memory_image()
-                    .clone(),
-            );
-        }
+        self.final_memory = Some(
+            self.chip_complex
+                .base
+                .memory_controller
+                .memory_image()
+                .clone(),
+        );
         #[cfg(feature = "bench-metrics")]
         self.finalize_metrics();