feat: simplify the handler type without Result

Author: jonathanpwang

crates/vm/derive/src/lib.rs

@@ -754,6 +754,13 @@ fn parse_executor_type(
 /// ```
 ///
 /// This will generate a TCO handler function with the same generics and where clauses.
+///
+/// # Safety
+///
+/// Do not use this macro if your function wants to terminate execution without error with a
+/// specific error code. The handler generated by this macro assumes that execution should continue
+/// unless the execute_impl returns an error. This is done for performance to skip an exit code
+/// check.
 #[proc_macro_attribute]
 pub fn create_tco_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
     tco::tco_impl(item)

crates/vm/derive/src/tco.rs

@@ -46,25 +46,26 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
                 ::openvm_circuit::system::memory::online::GuestMemory,
                 #ctx_type,
             >,
-        ) -> Result<(), ::openvm_circuit::arch::ExecutionError>
+        )
         #where_clause
         {
             let pre_compute = interpreter.get_pre_compute(exec_state.vm_state.pc);
             #execute_call;
 
-            if std::hint::unlikely(exec_state.exit_code.is_err()) {
-                return Err(::openvm_circuit::arch::ExecutionError::ExecStateError);
-            }
-            if std::hint::unlikely(exec_state.exit_code.as_ref().unwrap().is_some()) {
-                #ctx_type::on_terminate(exec_state);
-                // terminate
-                return Ok(());
+            if exec_state.exit_code.is_err() {
+                // stop execution
+                return;
             }
             if #ctx_type::should_suspend(exec_state) {
-                return Ok(());
+                return;
             }
             // exec_state.pc should have been updated by execute_impl at this point
-            let next_handler = interpreter.get_handler(exec_state.vm_state.pc)?;
+            let next_handler = interpreter.get_handler(exec_state.vm_state.pc);
+            if next_handler.is_none() {
+                exec_state.exit_code = Err(interpreter.pc_out_of_bounds_err(exec_state.vm_state.pc));
+                return;
+            }
+            let next_handler = next_handler.unwrap_unchecked();
 
             // The `become` keyword has a bug that is not re-passing the `interpreter`, `exec_state` references properly. But llvm seems to almost always guarantee tail call elimination when the function signature is the same as the current function.
             next_handler(interpreter, exec_state)

crates/vm/src/arch/execution.rs

@@ -109,7 +109,7 @@ pub type ExecuteFunc<F, CTX> =
 pub type Handler<F, CTX> = unsafe fn(
     interpreter: &InterpretedInstance<F, CTX>,
     exec_state: &mut VmExecState<F, GuestMemory, CTX>,
-) -> Result<(), ExecutionError>;
+);
 
 /// Trait for pure execution via a host interpreter. The trait methods provide the methods to
 /// pre-process the program code into function pointers which operate on `pre_compute` instruction

crates/vm/src/arch/interpreter.rs

@@ -189,21 +189,19 @@ where
         }
     }
 
+    pub fn pc_out_of_bounds_err(&self, pc: u32) -> ExecutionError {
+        ExecutionError::PcOutOfBounds {
+            pc,
+            pc_base: self.pc_base,
+            program_len: self.pre_compute_insns.len(),
+        }
+    }
+
     #[cfg(feature = "tco")]
     #[inline(always)]
-    pub fn get_handler(&self, pc: u32) -> Result<Handler<F, Ctx>, ExecutionError> {
+    pub fn get_handler(&self, pc: u32) -> Option<Handler<F, Ctx>> {
         let pc_idx = get_pc_index(self.pc_base, pc);
-        if std::hint::unlikely(pc_idx >= self.handlers.len()) {
-            return Err(ExecutionError::PcOutOfBounds {
-                pc,
-                pc_base: self.pc_base,
-                program_len: self.handlers.len(),
-            });
-        }
-        // SAFETY:
-        // - we checked above that pc_idx is within bounds
-        let handler = unsafe { self.handlers.get_unchecked(pc_idx) };
-        Ok(*handler)
+        self.handlers.get(pc_idx).map(|x| *x)
     }
 }
 
@@ -302,15 +300,21 @@ where
         }
         #[cfg(feature = "tco")]
         unsafe {
-            let handler = self.get_handler(exec_state.pc)?;
-            let res = handler(self, &mut exec_state);
-            if let Err(err) = res {
-                match err {
-                    ExecutionError::ExecStateError => {}
-                    _ => {
-                        return Err(err);
-                    }
-                }
+            let handler = self
+                .get_handler(exec_state.pc)
+                .ok_or(ExecutionError::PcOutOfBounds {
+                    pc: exec_state.pc,
+                    pc_base: self.pc_base,
+                    program_len: self.handlers.len(),
+                })?;
+            handler(self, &mut exec_state);
+
+            if exec_state
+                .exit_code
+                .as_ref()
+                .is_ok_and(|exit_code| exit_code.is_some())
+            {
+                ExecutionCtx::on_terminate(&mut exec_state);
             }
         }
 
@@ -540,17 +544,16 @@ unsafe fn terminate_execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
 unsafe fn terminate_execute_e12_tco_handler<F: PrimeField32, CTX: ExecutionCtxTrait>(
     interpreter: &InterpretedInstance<F, CTX>,
     vm_state: &mut VmExecState<F, GuestMemory, CTX>,
-) -> Result<(), ExecutionError> {
+) {
     let pre_compute = interpreter.get_pre_compute(vm_state.pc);
     terminate_execute_e12_impl(pre_compute, vm_state);
-    Ok(())
 }
 #[cfg(feature = "tco")]
 unsafe fn unreachable_tco_handler<F: PrimeField32, CTX>(
     _: &InterpretedInstance<F, CTX>,
     vm_state: &mut VmExecState<F, GuestMemory, CTX>,
-) -> Result<(), ExecutionError> {
-    Err(ExecutionError::Unreachable(vm_state.pc))
+) {
+    vm_state.exit_code = Err(ExecutionError::Unreachable(vm_state.pc));
 }
 
 fn get_pre_compute_max_size<F, E: Executor<F>>(

crates/vm/src/lib.rs

@@ -1,5 +1,4 @@
 #![cfg_attr(feature = "tco", allow(incomplete_features))]
-#![cfg_attr(feature = "tco", feature(likely_unlikely))]
 #![cfg_attr(feature = "tco", feature(explicit_tail_calls))]
 extern crate self as openvm_circuit;
 

crates/vm/src/system/phantom/execution.rs

@@ -101,7 +101,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
     vm_state.instret += 1;
 }
 
-#[cfg_attr(feature = "tco", create_tco_handler)]
+#[create_tco_handler]
 #[inline(always)]
 unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
     pre_compute: &[u8],

crates/vm/src/system/public_values/execution.rs

@@ -179,7 +179,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX, const B_IS_IMM: bool, const C_I
     state.instret += 1;
 }
 
-#[cfg_attr(feature = "tco", create_tco_handler)]
+#[create_tco_handler]
 #[inline(always)]
 unsafe fn execute_e1_impl<F: PrimeField32, CTX, const B_IS_IMM: bool, const C_IS_IMM: bool>(
     pre_compute: &[u8],

extensions/algebra/circuit/src/lib.rs

@@ -1,5 +1,4 @@
 #![cfg_attr(feature = "tco", allow(incomplete_features))]
-#![cfg_attr(feature = "tco", feature(likely_unlikely))]
 #![cfg_attr(feature = "tco", feature(explicit_tail_calls))]
 
 use derive_more::derive::{Deref, DerefMut};

extensions/bigint/circuit/src/lib.rs

@@ -1,5 +1,4 @@
 #![cfg_attr(feature = "tco", allow(incomplete_features))]
-#![cfg_attr(feature = "tco", feature(likely_unlikely))]
 #![cfg_attr(feature = "tco", feature(explicit_tail_calls))]
 use openvm_circuit::{
     self,

extensions/ecc/circuit/src/lib.rs

@@ -1,5 +1,4 @@
 #![cfg_attr(feature = "tco", allow(incomplete_features))]
-#![cfg_attr(feature = "tco", feature(likely_unlikely))]
 #![cfg_attr(feature = "tco", feature(explicit_tail_calls))]
 
 mod weierstrass_chip;