fix: RootVerifierLocalProver should respect AIR heights (#1554)

- Not a soundness issue. Just an issue of `RootVerifierLocalProver`
implementation in SDK.
- In some edge cases, the trace heights of the root verifier circuit are
less than the required heights. In this case, `RootVerifierLocalProver`
should pad its traces to the required heights.
- [Breaking change] Add a new field `internal_heights` into
`RootVerifierProvingKey`.

Author: nyunyunyunyu

crates/sdk/src/keygen/dummy.rs

@@ -63,9 +63,9 @@ pub(super) fn compute_root_proof_heights(
         .into_iter()
         .map(next_power_of_two_or_zero)
         .collect();
-    let mut internal_heights = res.internal_heights;
-    internal_heights.round_to_next_power_of_two_or_zero();
-    (air_heights, internal_heights)
+    let mut vm_heights = res.vm_heights;
+    vm_heights.round_to_next_power_of_two_or_zero();
+    (air_heights, vm_heights)
 }
 
 pub(super) fn dummy_internal_proof(
@@ -186,9 +186,9 @@ where
         assert_eq!(results.len(), 1, "dummy exe should have only 1 segment");
         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
+        let mut vm_heights = result.chip_complex.get_internal_trace_heights();
+        vm_heights.round_to_next_power_of_two();
+        vm_heights
     };
     // For the dummy proof, we must override the trace heights.
     let app_prover =

crates/sdk/src/keygen/mod.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use derivative::Derivative;
 use dummy::{compute_root_proof_heights, dummy_internal_proof_riscv_app_vm};
 use openvm_circuit::{
-    arch::{VirtualMachine, VmConfig},
+    arch::{VirtualMachine, VmComplexTraceHeights, VmConfig},
     system::{memory::dimensions::MemoryDimensions, program::trace::VmCommittedExe},
 };
 use openvm_continuations::{
@@ -333,7 +333,7 @@ impl AggStarkProvingKey {
             let mut vm_pk = vm.keygen();
             assert!(vm_pk.max_constraint_degree <= config.root_fri_params.max_constraint_degree());
 
-            let (air_heights, _internal_heights) = compute_root_proof_heights(
+            let (air_heights, vm_heights) = compute_root_proof_heights(
                 root_vm_config.clone(),
                 root_committed_exe.exe.clone(),
                 &internal_proof,
@@ -349,6 +349,7 @@ impl AggStarkProvingKey {
                 }),
                 root_committed_exe,
                 air_heights,
+                vm_heights,
             }
         };
         (
@@ -392,9 +393,8 @@ pub struct RootVerifierProvingKey {
     pub root_committed_exe: Arc<VmCommittedExe<RootSC>>,
     /// The constant trace heights, ordered by AIR ID.
     pub air_heights: Vec<usize>,
-    // The following is currently not used:
-    // The constant trace heights, ordered according to an internal ordering determined by the `NativeConfig`.
-    // pub internal_heights: VmComplexTraceHeights,
+    /// The constant trace heights in a semantic way for VM.
+    pub vm_heights: VmComplexTraceHeights,
 }
 
 impl RootVerifierProvingKey {

crates/sdk/src/prover/root.rs

@@ -52,7 +52,8 @@ impl RootVerifierLocalProver {
 impl SingleSegmentVmProver<RootSC> for RootVerifierLocalProver {
     fn prove(&self, input: impl Into<Streams<F>>) -> Proof<RootSC> {
         let input = input.into();
-        let vm = SingleSegmentVmExecutor::new(self.vm_config().clone());
+        let mut vm = SingleSegmentVmExecutor::new(self.vm_config().clone());
+        vm.set_override_trace_heights(self.root_verifier_pk.vm_heights.clone());
         let mut proof_input = vm
             .execute_and_generate(self.root_verifier_pk.root_committed_exe.clone(), input)
             .unwrap();

crates/vm/src/arch/vm.rs

@@ -376,7 +376,7 @@ pub struct SingleSegmentVmExecutionResult<F> {
     /// Heights of each AIR, ordered by AIR ID.
     pub air_heights: Vec<usize>,
     /// Heights of (SystemBase, Inventory), in an internal ordering.
-    pub internal_heights: VmComplexTraceHeights,
+    pub vm_heights: VmComplexTraceHeights,
 }
 
 impl<F, VC> SingleSegmentVmExecutor<F, VC>
@@ -424,7 +424,7 @@ where
             segment
         };
         let air_heights = segment.chip_complex.current_trace_heights();
-        let internal_heights = segment.chip_complex.get_internal_trace_heights();
+        let vm_heights = segment.chip_complex.get_internal_trace_heights();
         let public_values = if let Some(pv_chip) = segment.chip_complex.public_values_chip() {
             pv_chip.core.get_custom_public_values()
         } else {
@@ -433,7 +433,7 @@ where
         Ok(SingleSegmentVmExecutionResult {
             public_values,
             air_heights,
-            internal_heights,
+            vm_heights,
         })
     }
 

crates/vm/tests/integration_test.rs

@@ -132,7 +132,7 @@ fn test_vm_override_executor_height() {
         .unwrap();
     // Memory trace heights are not computed during execution.
     assert_eq!(
-        res.internal_heights.system,
+        res.vm_heights.system,
         SystemTraceHeights {
             memory: MemoryTraceHeights::Volatile(VolatileMemoryTraceHeights {
                 boundary: 1,
@@ -141,7 +141,7 @@ fn test_vm_override_executor_height() {
         }
     );
     assert_eq!(
-        res.internal_heights.inventory,
+        res.vm_heights.inventory,
         VmInventoryTraceHeights {
             chips: vec![
                 (ChipId::Executor(0), 0),