feat: add macro to generate tco handler and update interpreter for tco

Author: jonathanpwang

Cargo.toml

@@ -229,6 +229,7 @@ dashmap = "6.1.0"
 memmap2 = "0.9.5"
 libc = "0.2.175"
 tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }
+paste = "1.0.15"
 
 # default-features = false for no_std for use in guest programs
 itertools = { version = "0.14.0", default-features = false }

crates/vm/Cargo.toml

@@ -50,7 +50,7 @@ openvm-native-compiler.workspace = true
 openvm-rv32im-transpiler.workspace = true
 
 [features]
-default = ["parallel", "jemalloc"]
+default = ["parallel", "jemalloc", "tco"]
 parallel = [
     "openvm-stark-backend/parallel",
     "dashmap/rayon",

crates/vm/derive/Cargo.toml

@@ -10,7 +10,7 @@ license.workspace = true
 proc-macro = true
 
 [dependencies]
-syn = { version = "2.0", features = ["parsing"] }
+syn = { version = "2.0", features = ["parsing", "full"] }
 quote = "1.0"
 proc-macro2 = "1.0"
 itertools = { workspace = true }

crates/vm/derive/src/lib.rs

@@ -9,6 +9,8 @@ use syn::{
     GenericParam, Ident, Meta, Token,
 };
 
+mod tco;
+
 #[proc_macro_derive(PreflightExecutor)]
 pub fn preflight_executor_derive(input: TokenStream) -> TokenStream {
     let ast: syn::DeriveInput = syn::parse(input).unwrap();
@@ -172,6 +174,18 @@ pub fn executor_derive(input: TokenStream) -> TokenStream {
                         Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
                         self.0.pre_compute(pc, inst, data)
                     }
+
+                    #[cfg(feature = "tco")]
+                    fn handler<Ctx>(
+                        &self,
+                        pc: u32,
+                        inst: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
+                        data: &mut [u8],
+                    ) -> Result<::openvm_circuit::arch::Handler<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
+                    where
+                        Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
+                        self.0.handler(pc, inst, data)
+                    }
                 }
             }
             .into()
@@ -205,18 +219,21 @@ pub fn executor_derive(input: TokenStream) -> TokenStream {
                 });
             // Use full path ::openvm_circuit... so it can be used either within or outside the vm
             // crate. Assume F is already generic of the field.
-            let (pre_compute_size_arms, pre_compute_arms, where_predicates): (Vec<_>, Vec<_>, Vec<_>) = multiunzip(variants.iter().map(|(variant_name, field)| {
+            let (pre_compute_size_arms, pre_compute_arms, handler_arms, where_predicates): (Vec<_>, Vec<_>, Vec<_>, Vec<_>) = multiunzip(variants.iter().map(|(variant_name, field)| {
                 let field_ty = &field.ty;
                 let pre_compute_size_arm = quote! {
                     #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::Executor<#first_ty_generic>>::pre_compute_size(x)
                 };
                 let pre_compute_arm = quote! {
                     #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::Executor<#first_ty_generic>>::pre_compute(x, pc, instruction, data)
                 };
+                let handler_arm = quote! {
+                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::Executor<#first_ty_generic>>::handler(x, pc, instruction, data)
+                };
                 let where_predicate = syn::parse_quote! {
                     #field_ty: ::openvm_circuit::arch::Executor<#first_ty_generic>
                 };
-                (pre_compute_size_arm, pre_compute_arm, where_predicate)
+                (pre_compute_size_arm, pre_compute_arm, handler_arm, where_predicate)
             }));
             let where_clause = new_generics.make_where_clause();
             for predicate in where_predicates {
@@ -247,6 +264,20 @@ pub fn executor_derive(input: TokenStream) -> TokenStream {
                             #(#pre_compute_arms,)*
                         }
                     }
+
+                    #[cfg(feature = "tco")]
+                    fn handler<Ctx>(
+                        &self,
+                        pc: u32,
+                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
+                        data: &mut [u8],
+                    ) -> Result<::openvm_circuit::arch::Handler<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
+                    where
+                        Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
+                        match self {
+                            #(#handler_arms,)*
+                        }
+                    }
                 }
             }
             .into()
@@ -501,7 +532,7 @@ fn generate_config_traits_impl(name: &Ident, inner: &DataStruct) -> syn::Result<
         .iter()
         .filter(|f| f.attrs.iter().any(|attr| attr.path().is_ident("config")))
         .exactly_one()
-        .clone()
+        .ok()
         .expect("Exactly one field must have the #[config] attribute");
     let (source_name, source_name_upper) =
         gen_name_with_uppercase_idents(source_field.ident.as_ref().unwrap());
@@ -700,3 +731,30 @@ fn parse_executor_type(
         })
     }
 }
+
+/// An attribute procedural macro for creating TCO (Tail Call Optimization) handlers.
+///
+/// This macro generates a handler function that wraps an execute implementation
+/// with tail call optimization using the `become` keyword. It extracts the generics
+/// and where clauses from the original function.
+///
+/// # Usage
+///
+/// Place this attribute above a function definition:
+/// ```
+/// #[create_tco_handler = "handler_name"]
+/// unsafe fn execute_e1_impl<F: PrimeField32, CTX, const B_IS_IMM: bool>(
+///     pre_compute: &[u8],
+///     state: &mut VmExecState<F, GuestMemory, CTX>,
+/// ) where
+///     CTX: ExecutionCtxTrait,
+/// {
+///     // function body
+/// }
+/// ```
+///
+/// This will generate a TCO handler function with the same generics and where clauses.
+#[proc_macro_attribute]
+pub fn create_tco_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    tco::tco_impl(item)
+}

crates/vm/derive/src/tco.rs

@@ -0,0 +1,118 @@
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, ItemFn};
+
+/// Implementation of the TCO handler generation logic.
+/// This is called from the proc macro attribute in lib.rs.
+pub fn tco_impl(item: TokenStream) -> TokenStream {
+    // Parse the input function
+    let input_fn = parse_macro_input!(item as ItemFn);
+
+    // Extract information from the function
+    let fn_name = &input_fn.sig.ident;
+    let generics = &input_fn.sig.generics;
+    let where_clause = &generics.where_clause;
+
+    // Extract the first two generic type parameters (F and CTX)
+    let (f_type, ctx_type) = extract_f_and_ctx_types(generics);
+    // Derive new function name:
+    // If original ends with `_impl`, replace with `_tco_handler`, else append suffix.
+    let new_name_str = fn_name
+        .to_string()
+        .strip_suffix("_impl")
+        .map(|base| format!("{base}_tco_handler"))
+        .unwrap_or_else(|| format!("{fn_name}_tco_handler"));
+    let handler_name = format_ident!("{}", new_name_str);
+
+    // Build the generic parameters for the handler, preserving all original generics
+    let handler_generics = generics.clone();
+
+    // Build the function call with all the generics
+    let generic_args = build_generic_args(generics);
+    let execute_call = if generic_args.is_empty() {
+        quote! { #fn_name(pre_compute, exec_state) }
+    } else {
+        quote! { #fn_name::<#(#generic_args),*>(pre_compute, exec_state) }
+    };
+
+    // Generate the TCO handler function
+    let handler_fn = quote! {
+        #[cfg(feature = "tco")]
+        #[inline(never)]
+        unsafe fn #handler_name #handler_generics (
+            interpreter: &::openvm_circuit::arch::interpreter::InterpretedInstance<#f_type, #ctx_type>,
+            exec_state: &mut ::openvm_circuit::arch::VmExecState<
+                #f_type,
+                ::openvm_circuit::system::memory::online::GuestMemory,
+                #ctx_type,
+            >,
+        ) -> Result<(), ::openvm_circuit::arch::ExecutionError>
+        #where_clause
+        {
+            let pre_compute = interpreter.get_pre_compute(exec_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()) {
+                // terminate
+                return Ok(());
+            }
+            // exec_state.pc should have been updated by execute_impl at this point
+            let next_handler = interpreter.get_handler(exec_state.pc)?;
+            become next_handler(interpreter, exec_state)
+        }
+    };
+
+    // Return both the original function and the new handler
+    let output = quote! {
+        #input_fn
+
+        #handler_fn
+    };
+
+    TokenStream::from(output)
+}
+
+fn extract_f_and_ctx_types(generics: &syn::Generics) -> (syn::Ident, syn::Ident) {
+    let mut type_params = generics.params.iter().filter_map(|param| {
+        if let syn::GenericParam::Type(type_param) = param {
+            Some(&type_param.ident)
+        } else {
+            None
+        }
+    });
+
+    let f_type = type_params
+        .next()
+        .expect("Function must have at least one type parameter (F)")
+        .clone();
+    let ctx_type = type_params
+        .next()
+        .expect("Function must have at least two type parameters (F and CTX)")
+        .clone();
+
+    (f_type, ctx_type)
+}
+
+fn build_generic_args(generics: &syn::Generics) -> Vec<proc_macro2::TokenStream> {
+    generics
+        .params
+        .iter()
+        .map(|param| match param {
+            syn::GenericParam::Type(type_param) => {
+                let ident = &type_param.ident;
+                quote! { #ident }
+            }
+            syn::GenericParam::Lifetime(lifetime) => {
+                let lifetime = &lifetime.lifetime;
+                quote! { #lifetime }
+            }
+            syn::GenericParam::Const(const_param) => {
+                let ident = &const_param.ident;
+                quote! { #ident }
+            }
+        })
+        .collect()
+}

crates/vm/src/arch/execution.rs

@@ -12,6 +12,8 @@ use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use super::{execution_mode::ExecutionCtxTrait, Streams, VmExecState};
+#[cfg(feature = "tco")]
+use crate::arch::interpreter::InterpretedInstance;
 #[cfg(feature = "metrics")]
 use crate::metrics::VmMetrics;
 use crate::{
@@ -72,6 +74,9 @@ pub enum ExecutionError {
     Inventory(#[from] ExecutorInventoryError),
     #[error("static program error: {0}")]
     Static(#[from] StaticProgramError),
+    // Placeholder error type for tco
+    #[error("error in VmExecState")]
+    ExecStateError,
 }
 
 /// Errors in the program that can be statically analyzed before runtime.
@@ -91,7 +96,20 @@ pub enum StaticProgramError {
 /// The `pre_compute: &[u8]` is a pre-computed buffer of data corresponding to a single instruction.
 /// The contents of `pre_compute` are determined from the program code as specified by the
 /// [Executor] and [MeteredExecutor] traits.
-pub type ExecuteFunc<F, CTX> = unsafe fn(&[u8], &mut VmExecState<F, GuestMemory, CTX>);
+pub type ExecuteFunc<F, CTX> =
+    unsafe fn(pre_compute: &[u8], exec_state: &mut VmExecState<F, GuestMemory, CTX>);
+
+/// Handler for tail call elimination. The `CTX` is assumed to contain pointers to the pre-computed
+/// buffer and the function handler table.
+///
+/// - `pre_compute_buf` is the starting pointer of the pre-computed buffer.
+/// - `handlers` is the starting pointer of the table of function pointers of `Handler` type. The
+///   pointer is typeless to avoid self-referential types.
+#[cfg(feature = "tco")]
+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
@@ -108,6 +126,20 @@ pub trait Executor<F> {
     ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
     where
         Ctx: ExecutionCtxTrait;
+
+    /// Returns a function pointer with tail call optimization. The handler function assumes that
+    /// the pre-compute buffer it receives is the populated `data`.
+    // NOTE: we could have used `pre_compute` above to populate `data`, but the implementations were
+    // simpler to keep `handler` entirely separate from `pre_compute`.
+    #[cfg(feature = "tco")]
+    fn handler<Ctx>(
+        &self,
+        pc: u32,
+        inst: &Instruction<F>,
+        data: &mut [u8],
+    ) -> Result<Handler<F, Ctx>, StaticProgramError>
+    where
+        Ctx: ExecutionCtxTrait;
 }
 
 /// Trait for metered execution via a host interpreter. The trait methods provide the methods to

crates/vm/src/arch/execution_mode/pure.rs

@@ -19,12 +19,6 @@ impl ExecutionCtx {
     }
 }
 
-impl Default for ExecutionCtx {
-    fn default() -> Self {
-        Self::new(None)
-    }
-}
-
 impl ExecutionCtxTrait for ExecutionCtx {
     #[inline(always)]
     fn on_memory_operation(&mut self, _address_space: u32, _ptr: u32, _size: u32) {}