feat: updated stark_utils to test metered_cost and aot_metered_cost

Author: Maillew

crates/vm/src/arch/vm.rs

@@ -309,16 +309,15 @@ where
         InterpretedInstance::new_metered(&self.inventory, exe, executor_idx_to_air_idx)
     }
 
-    // Creates an instance of the AotInstance specialized for metered execution of the given
-    /// `exe`.
-    pub fn aot_metered_instance(
+    // Creates an AOT instance for metered execution of the given `exe`.
+    pub fn metered_cost_aot_instance(
         &self,
         exe: &VmExe<F>,
         executor_idx_to_air_idx: &[usize],
-    ) -> Result<AotInstance<F, MeteredCtx>, StaticProgramError> {
+    ) -> Result<AotInstance<F, MeteredCostCtx>, StaticProgramError> {
         #[cfg(all(target_arch = "x86_64", not(feature = "tco")))]
         {
-            AotInstance::new_metered(&self.inventory, exe, executor_idx_to_air_idx)
+            AotInstance::new_metered_cost(&self.inventory, exe, executor_idx_to_air_idx)
         }
         #[cfg(any(not(target_arch = "x86_64"), feature = "tco"))]
         {
@@ -543,6 +542,20 @@ where
             .metered_cost_instance(exe, &executor_idx_to_air_idx)
     }
 
+    // Metered Cost AOT execution
+    pub fn get_metered_cost_aot_instance(
+        &self,
+        exe: &VmExe<Val<E::SC>>,
+    ) -> Result<AotInstance<Val<E::SC>, MeteredCostCtx>, StaticProgramError>
+    where
+        Val<E::SC>: PrimeField32,
+        <VB::VmConfig as VmExecutionConfig<Val<E::SC>>>::Executor: MeteredExecutor<Val<E::SC>>,
+    {
+        let executor_idx_to_air_idx = self.executor_idx_to_air_idx();
+        self.executor()
+            .metered_cost_aot_instance(exe, &executor_idx_to_air_idx)
+    }
+
     pub fn preflight_interpreter(
         &self,
         exe: &VmExe<Val<E::SC>>,

crates/vm/src/utils/stark_utils.rs

@@ -115,7 +115,7 @@ where
     let exe = exe.into();
     let input = input.into();
     let metered_ctx = vm.build_metered_ctx(&exe);
-
+    let metered_cost_ctx = vm.build_metered_cost_ctx();
     /*
     Assertions for Pure Execution AOT
     */
@@ -174,7 +174,21 @@ where
             assert_eq!(segments[i].trace_heights, aot_segments[i].trace_heights);
         }
     }
+    /*
+    Assertions for Metered AOT Cost
+    */
+    {
+        let (cost, vm_state) = vm
+            .metered_cost_interpreter(&exe)?
+            .execute_metered_cost(input.clone(), metered_cost_ctx.clone())?;
+
+        let (aot_cost, aot_vm_state) = vm
+            .get_metered_cost_aot_instance(&exe)?
+            .execute_metered_cost(input.clone(), metered_cost_ctx.clone())?;
 
+        assert_eq!(vm_state.instret(), aot_vm_state.instret());
+        assert_eq!(cost, aot_cost);
+    }
     /* TODO: this is a temporary change to use `get_metered_aot_instance` instead of `metered_interpreter`
     to test AOT segments in addition to the equal assertions
     We would want to revert `stark_utils.rs` back to how it looked like in main