openvm_circuit_derive/
lib.rs

1extern crate alloc;
2extern crate proc_macro;
3
4use itertools::{multiunzip, Itertools};
5use proc_macro::{Span, TokenStream};
6use quote::{quote, ToTokens};
7use syn::{
8    parse_quote, punctuated::Punctuated, spanned::Spanned, Data, DataStruct, Field, Fields,
9    GenericParam, Ident, Meta, Token,
10};
11
12#[cfg(feature = "tco")]
13mod tco;
14
15#[proc_macro_derive(PreflightExecutor)]
16pub fn preflight_executor_derive(input: TokenStream) -> TokenStream {
17    let ast: syn::DeriveInput = syn::parse(input).unwrap();
18
19    let name = &ast.ident;
20    let generics = &ast.generics;
21    let (_, ty_generics, _) = generics.split_for_impl();
22
23    let default_ty_generic = Ident::new("F", proc_macro2::Span::call_site());
24    let mut new_generics = generics.clone();
25    new_generics.params.push(syn::parse_quote! { RA });
26    let field_ty_generic = generics
27        .params
28        .first()
29        .and_then(|param| match param {
30            GenericParam::Type(type_param) => Some(&type_param.ident),
31            _ => None,
32        })
33        .unwrap_or_else(|| {
34            new_generics.params.push(syn::parse_quote! { F });
35            &default_ty_generic
36        });
37
38    match &ast.data {
39        Data::Struct(inner) => {
40            // Check if the struct has only one unnamed field
41            let inner_ty = match &inner.fields {
42                Fields::Unnamed(fields) => {
43                    if fields.unnamed.len() != 1 {
44                        panic!("Only one unnamed field is supported");
45                    }
46                    fields.unnamed.first().unwrap().ty.clone()
47                }
48                _ => panic!("Only unnamed fields are supported"),
49            };
50            // Use full path ::openvm_circuit... so it can be used either within or outside the vm
51            // crate.
52            let where_clause = new_generics.make_where_clause();
53            where_clause.predicates.push(
54                syn::parse_quote! { #inner_ty: ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA> },
55            );
56            let (impl_generics, _, where_clause) = new_generics.split_for_impl();
57            quote! {
58                impl #impl_generics ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA> for #name #ty_generics #where_clause {
59                    fn execute(
60                        &self,
61                        state: ::openvm_circuit::arch::VmStateMut<#field_ty_generic, ::openvm_circuit::system::memory::online::TracingMemory, RA>,
62                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<#field_ty_generic>,
63                    ) -> Result<(), ::openvm_circuit::arch::ExecutionError> {
64                        self.0.execute(state, instruction)
65                    }
66
67                    fn get_opcode_name(&self, opcode: usize) -> String {
68                        self.0.get_opcode_name(opcode)
69                    }
70                }
71            }
72            .into()
73        }
74        Data::Enum(e) => {
75            let variants = e
76                .variants
77                .iter()
78                .map(|variant| {
79                    let variant_name = &variant.ident;
80
81                    let mut fields = variant.fields.iter();
82                    let field = fields.next().unwrap();
83                    assert!(fields.next().is_none(), "Only one field is supported");
84                    (variant_name, field)
85                })
86                .collect::<Vec<_>>();
87            // Use full path ::openvm_circuit... so it can be used either within or outside the vm
88            // crate. Assume F is already generic of the field.
89            let (execute_arms, get_opcode_name_arms, where_predicates): (Vec<_>, Vec<_>, Vec<_>) =
90                multiunzip(variants.iter().map(|(variant_name, field)| {
91                    let field_ty = &field.ty;
92                    let execute_arm = quote! {
93                        #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA>>::execute(x, state, instruction)
94                    };
95                    let get_opcode_name_arm = quote! {
96                        #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA>>::get_opcode_name(x, opcode)
97                    };
98                    let where_predicate = syn::parse_quote! {
99                        #field_ty: ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA>
100                    };
101                    (execute_arm, get_opcode_name_arm, where_predicate)
102                }));
103            let where_clause = new_generics.make_where_clause();
104            for predicate in where_predicates {
105                where_clause.predicates.push(predicate);
106            }
107            // Don't use these ty_generics because it might have extra "F"
108            let (impl_generics, _, where_clause) = new_generics.split_for_impl();
109            quote! {
110                impl #impl_generics ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA> for #name #ty_generics #where_clause {
111                    fn execute(
112                        &self,
113                        state: ::openvm_circuit::arch::VmStateMut<#field_ty_generic, ::openvm_circuit::system::memory::online::TracingMemory, RA>,
114                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<#field_ty_generic>,
115                    ) -> Result<(), ::openvm_circuit::arch::ExecutionError> {
116                        match self {
117                            #(#execute_arms,)*
118                        }
119                    }
120
121                    fn get_opcode_name(&self, opcode: usize) -> String {
122                        match self {
123                            #(#get_opcode_name_arms,)*
124                        }
125                    }
126                }
127            }
128            .into()
129        }
130        Data::Union(_) => unimplemented!("Unions are not supported"),
131    }
132}
133
134#[proc_macro_derive(Executor)]
135pub fn executor_derive(input: TokenStream) -> TokenStream {
136    let ast: syn::DeriveInput = syn::parse(input).unwrap();
137
138    let name = &ast.ident;
139    let generics = &ast.generics;
140    let (impl_generics, ty_generics, _) = generics.split_for_impl();
141
142    match &ast.data {
143        Data::Struct(inner) => {
144            // Check if the struct has only one unnamed field
145            let inner_ty = match &inner.fields {
146                Fields::Unnamed(fields) => {
147                    if fields.unnamed.len() != 1 {
148                        panic!("Only one unnamed field is supported");
149                    }
150                    fields.unnamed.first().unwrap().ty.clone()
151                }
152                _ => panic!("Only unnamed fields are supported"),
153            };
154            // Use full path ::openvm_circuit... so it can be used either within or outside the vm
155            // crate. Assume F is already generic of the field.
156            let mut new_generics = generics.clone();
157            let where_clause = new_generics.make_where_clause();
158            where_clause
159                .predicates
160                .push(syn::parse_quote! { #inner_ty: ::openvm_circuit::arch::Executor<F> });
161
162            // We use the macro's feature to decide whether to generate the impl or not. This avoids
163            // the target crate needing the "tco" feature defined.
164            #[cfg(feature = "tco")]
165            let handler = quote! {
166                fn handler<Ctx>(
167                    &self,
168                    pc: u32,
169                    inst: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
170                    data: &mut [u8],
171                ) -> Result<::openvm_circuit::arch::Handler<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
172                where
173                    Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
174                    self.0.handler(pc, inst, data)
175                }
176            };
177            #[cfg(not(feature = "tco"))]
178            let handler = quote! {};
179
180            quote! {
181                impl #impl_generics ::openvm_circuit::arch::Executor<F> for #name #ty_generics #where_clause {
182                    #[inline(always)]
183                    fn pre_compute_size(&self) -> usize {
184                        self.0.pre_compute_size()
185                    }
186                    #[inline(always)]
187                    fn pre_compute<Ctx>(
188                        &self,
189                        pc: u32,
190                        inst: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
191                        data: &mut [u8],
192                    ) -> Result<::openvm_circuit::arch::ExecuteFunc<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
193                    where
194                        Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
195                        self.0.pre_compute(pc, inst, data)
196                    }
197
198                    #handler
199                }
200            }
201            .into()
202        }
203        Data::Enum(e) => {
204            let variants = e
205                .variants
206                .iter()
207                .map(|variant| {
208                    let variant_name = &variant.ident;
209
210                    let mut fields = variant.fields.iter();
211                    let field = fields.next().unwrap();
212                    assert!(fields.next().is_none(), "Only one field is supported");
213                    (variant_name, field)
214                })
215                .collect::<Vec<_>>();
216            let default_ty_generic = Ident::new("F", proc_macro2::Span::call_site());
217            let mut new_generics = generics.clone();
218            let first_ty_generic = ast
219                .generics
220                .params
221                .first()
222                .and_then(|param| match param {
223                    GenericParam::Type(type_param) => Some(&type_param.ident),
224                    _ => None,
225                })
226                .unwrap_or_else(|| {
227                    new_generics.params.push(syn::parse_quote! { F });
228                    &default_ty_generic
229                });
230            // Use full path ::openvm_circuit... so it can be used either within or outside the vm
231            // crate. Assume F is already generic of the field.
232            let (pre_compute_size_arms, pre_compute_arms, _handler_arms, where_predicates): (Vec<_>, Vec<_>, Vec<_>, Vec<_>) = multiunzip(variants.iter().map(|(variant_name, field)| {
233                let field_ty = &field.ty;
234                let pre_compute_size_arm = quote! {
235                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::Executor<#first_ty_generic>>::pre_compute_size(x)
236                };
237                let pre_compute_arm = quote! {
238                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::Executor<#first_ty_generic>>::pre_compute(x, pc, instruction, data)
239                };
240                let handler_arm = quote! {
241                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::Executor<#first_ty_generic>>::handler(x, pc, instruction, data)
242                };
243                let where_predicate = syn::parse_quote! {
244                    #field_ty: ::openvm_circuit::arch::Executor<#first_ty_generic>
245                };
246                (pre_compute_size_arm, pre_compute_arm, handler_arm, where_predicate)
247            }));
248            let where_clause = new_generics.make_where_clause();
249            for predicate in where_predicates {
250                where_clause.predicates.push(predicate);
251            }
252            // We use the macro's feature to decide whether to generate the impl or not. This avoids
253            // the target crate needing the "tco" feature defined.
254            #[cfg(feature = "tco")]
255            let handler = quote! {
256                fn handler<Ctx>(
257                    &self,
258                    pc: u32,
259                    instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
260                    data: &mut [u8],
261                ) -> Result<::openvm_circuit::arch::Handler<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
262                where
263                    Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
264                    match self {
265                        #(#_handler_arms,)*
266                    }
267                }
268            };
269            #[cfg(not(feature = "tco"))]
270            let handler = quote! {};
271
272            // Don't use these ty_generics because it might have extra "F"
273            let (impl_generics, _, where_clause) = new_generics.split_for_impl();
274
275            quote! {
276                impl #impl_generics ::openvm_circuit::arch::Executor<#first_ty_generic> for #name #ty_generics #where_clause {
277                    #[inline(always)]
278                    fn pre_compute_size(&self) -> usize {
279                        match self {
280                            #(#pre_compute_size_arms,)*
281                        }
282                    }
283
284                    #[inline(always)]
285                    fn pre_compute<Ctx>(
286                        &self,
287                        pc: u32,
288                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
289                        data: &mut [u8],
290                    ) -> Result<::openvm_circuit::arch::ExecuteFunc<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
291                    where
292                        Ctx: ::openvm_circuit::arch::execution_mode::ExecutionCtxTrait, {
293                        match self {
294                            #(#pre_compute_arms,)*
295                        }
296                    }
297
298                    #handler
299                }
300            }
301            .into()
302        }
303        Data::Union(_) => unimplemented!("Unions are not supported"),
304    }
305}
306
307#[proc_macro_derive(MeteredExecutor)]
308pub fn metered_executor_derive(input: TokenStream) -> TokenStream {
309    let ast: syn::DeriveInput = syn::parse(input).unwrap();
310
311    let name = &ast.ident;
312    let generics = &ast.generics;
313    let (impl_generics, ty_generics, _) = generics.split_for_impl();
314
315    match &ast.data {
316        Data::Struct(inner) => {
317            // Check if the struct has only one unnamed field
318            let inner_ty = match &inner.fields {
319                Fields::Unnamed(fields) => {
320                    if fields.unnamed.len() != 1 {
321                        panic!("Only one unnamed field is supported");
322                    }
323                    fields.unnamed.first().unwrap().ty.clone()
324                }
325                _ => panic!("Only unnamed fields are supported"),
326            };
327            // Use full path ::openvm_circuit... so it can be used either within or outside the vm
328            // crate. Assume F is already generic of the field.
329            let mut new_generics = generics.clone();
330            let where_clause = new_generics.make_where_clause();
331            where_clause
332                .predicates
333                .push(syn::parse_quote! { #inner_ty: ::openvm_circuit::arch::MeteredExecutor<F> });
334
335            // We use the macro's feature to decide whether to generate the impl or not. This avoids
336            // the target crate needing the "tco" feature defined.
337            #[cfg(feature = "tco")]
338            let metered_handler = quote! {
339                fn metered_handler<Ctx>(
340                    &self,
341                    chip_idx: usize,
342                    pc: u32,
343                    inst: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
344                    data: &mut [u8],
345                ) -> Result<::openvm_circuit::arch::Handler<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
346                where
347                    Ctx: ::openvm_circuit::arch::execution_mode::MeteredExecutionCtxTrait, {
348                    self.0.metered_handler(chip_idx, pc, inst, data)
349                }
350            };
351            #[cfg(not(feature = "tco"))]
352            let metered_handler = quote! {};
353
354            quote! {
355                impl #impl_generics ::openvm_circuit::arch::MeteredExecutor<F> for #name #ty_generics #where_clause {
356                    #[inline(always)]
357                    fn metered_pre_compute_size(&self) -> usize {
358                        self.0.metered_pre_compute_size()
359                    }
360                    #[inline(always)]
361                    fn metered_pre_compute<Ctx>(
362                        &self,
363                        chip_idx: usize,
364                        pc: u32,
365                        inst: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
366                        data: &mut [u8],
367                    ) -> Result<::openvm_circuit::arch::ExecuteFunc<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
368                    where
369                        Ctx: ::openvm_circuit::arch::execution_mode::MeteredExecutionCtxTrait, {
370                        self.0.metered_pre_compute(chip_idx, pc, inst, data)
371                    }
372                    #metered_handler
373                }
374            }
375                .into()
376        }
377        Data::Enum(e) => {
378            let variants = e
379                .variants
380                .iter()
381                .map(|variant| {
382                    let variant_name = &variant.ident;
383
384                    let mut fields = variant.fields.iter();
385                    let field = fields.next().unwrap();
386                    assert!(fields.next().is_none(), "Only one field is supported");
387                    (variant_name, field)
388                })
389                .collect::<Vec<_>>();
390            let default_ty_generic = Ident::new("F", proc_macro2::Span::call_site());
391            let mut new_generics = generics.clone();
392            let first_ty_generic = ast
393                .generics
394                .params
395                .first()
396                .and_then(|param| match param {
397                    GenericParam::Type(type_param) => Some(&type_param.ident),
398                    _ => None,
399                })
400                .unwrap_or_else(|| {
401                    new_generics.params.push(syn::parse_quote! { F });
402                    &default_ty_generic
403                });
404            // Use full path ::openvm_circuit... so it can be used either within or outside the vm
405            // crate. Assume F is already generic of the field.
406            let (pre_compute_size_arms, metered_pre_compute_arms, _metered_handler_arms, where_predicates): (Vec<_>, Vec<_>, Vec<_>, Vec<_>) = multiunzip(variants.iter().map(|(variant_name, field)| {
407                let field_ty = &field.ty;
408                let pre_compute_size_arm = quote! {
409                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::MeteredExecutor<#first_ty_generic>>::metered_pre_compute_size(x)
410                };
411                let metered_pre_compute_arm = quote! {
412                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::MeteredExecutor<#first_ty_generic>>::metered_pre_compute(x, chip_idx, pc, instruction, data)
413                };
414                let metered_handler_arm = quote! {
415                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::MeteredExecutor<#first_ty_generic>>::metered_handler(x, chip_idx, pc, instruction, data)
416                };
417                let where_predicate = syn::parse_quote! {
418                    #field_ty: ::openvm_circuit::arch::MeteredExecutor<#first_ty_generic>
419                };
420                (pre_compute_size_arm, metered_pre_compute_arm, metered_handler_arm, where_predicate)
421            }));
422            let where_clause = new_generics.make_where_clause();
423            for predicate in where_predicates {
424                where_clause.predicates.push(predicate);
425            }
426            // Don't use these ty_generics because it might have extra "F"
427            let (impl_generics, _, where_clause) = new_generics.split_for_impl();
428
429            // We use the macro's feature to decide whether to generate the impl or not. This avoids
430            // the target crate needing the "tco" feature defined.
431            #[cfg(feature = "tco")]
432            let metered_handler = quote! {
433                fn metered_handler<Ctx>(
434                    &self,
435                    chip_idx: usize,
436                    pc: u32,
437                    instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
438                    data: &mut [u8],
439                ) -> Result<::openvm_circuit::arch::Handler<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
440                where
441                    Ctx: ::openvm_circuit::arch::execution_mode::MeteredExecutionCtxTrait,
442                {
443                    match self {
444                        #(#_metered_handler_arms,)*
445                    }
446                }
447            };
448            #[cfg(not(feature = "tco"))]
449            let metered_handler = quote! {};
450
451            quote! {
452                impl #impl_generics ::openvm_circuit::arch::MeteredExecutor<#first_ty_generic> for #name #ty_generics #where_clause {
453                    #[inline(always)]
454                    fn metered_pre_compute_size(&self) -> usize {
455                        match self {
456                            #(#pre_compute_size_arms,)*
457                        }
458                    }
459
460                    #[inline(always)]
461                    fn metered_pre_compute<Ctx>(
462                        &self,
463                        chip_idx: usize,
464                        pc: u32,
465                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
466                        data: &mut [u8],
467                    ) -> Result<::openvm_circuit::arch::ExecuteFunc<F, Ctx>, ::openvm_circuit::arch::StaticProgramError>
468                    where
469                        Ctx: ::openvm_circuit::arch::execution_mode::MeteredExecutionCtxTrait, {
470                        match self {
471                            #(#metered_pre_compute_arms,)*
472                        }
473                    }
474
475                    #metered_handler
476                }
477            }
478                .into()
479        }
480        Data::Union(_) => unimplemented!("Unions are not supported"),
481    }
482}
483
484/// Derives `AnyEnum` trait on an enum type.
485/// By default an enum arm will just return `self` as `&dyn Any`.
486///
487/// Use the `#[any_enum]` field attribute to specify that the
488/// arm itself implements `AnyEnum` and should call the inner `as_any_kind` method.
489#[proc_macro_derive(AnyEnum, attributes(any_enum))]
490pub fn any_enum_derive(input: TokenStream) -> TokenStream {
491    let ast: syn::DeriveInput = syn::parse(input).unwrap();
492
493    let name = &ast.ident;
494    let generics = &ast.generics;
495    let (impl_generics, ty_generics, _) = generics.split_for_impl();
496
497    match &ast.data {
498        Data::Enum(e) => {
499            let variants = e
500                .variants
501                .iter()
502                .map(|variant| {
503                    let variant_name = &variant.ident;
504
505                    // Check if the variant has #[any_enum] attribute
506                    let is_enum = variant
507                        .attrs
508                        .iter()
509                        .any(|attr| attr.path().is_ident("any_enum"));
510                    let mut fields = variant.fields.iter();
511                    let field = fields.next().unwrap();
512                    assert!(fields.next().is_none(), "Only one field is supported");
513                    (variant_name, field, is_enum)
514                })
515                .collect::<Vec<_>>();
516            let (arms, arms_mut): (Vec<_>, Vec<_>) =
517                variants.iter().map(|(variant_name, field, is_enum)| {
518                    let field_ty = &field.ty;
519
520                    if *is_enum {
521                        // Call the inner trait impl
522                        (quote! {
523                            #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::AnyEnum>::as_any_kind(x)
524                        },
525                        quote! {
526                            #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::AnyEnum>::as_any_kind_mut(x)
527                        })
528                    } else {
529                        (quote! {
530                            #name::#variant_name(x) => x
531                        },
532                        quote! {
533                            #name::#variant_name(x) => x
534                        })
535                    }
536                }).unzip();
537            quote! {
538                impl #impl_generics ::openvm_circuit::arch::AnyEnum for #name #ty_generics {
539                    fn as_any_kind(&self) -> &dyn std::any::Any {
540                        match self {
541                            #(#arms,)*
542                        }
543                    }
544
545                    fn as_any_kind_mut(&mut self) -> &mut dyn std::any::Any {
546                        match self {
547                            #(#arms_mut,)*
548                        }
549                    }
550                }
551            }
552            .into()
553        }
554        _ => syn::Error::new(name.span(), "Only enums are supported")
555            .to_compile_error()
556            .into(),
557    }
558}
559
560#[proc_macro_derive(VmConfig, attributes(config, extension))]
561pub fn vm_generic_config_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
562    let ast = syn::parse_macro_input!(input as syn::DeriveInput);
563    let name = &ast.ident;
564
565    match &ast.data {
566        syn::Data::Struct(inner) => match generate_config_traits_impl(name, inner) {
567            Ok(tokens) => tokens,
568            Err(err) => err.to_compile_error().into(),
569        },
570        _ => syn::Error::new(name.span(), "Only structs are supported")
571            .to_compile_error()
572            .into(),
573    }
574}
575
576fn generate_config_traits_impl(name: &Ident, inner: &DataStruct) -> syn::Result<TokenStream> {
577    let gen_name_with_uppercase_idents = |ident: &Ident| {
578        let mut name = ident.to_string().chars().collect::<Vec<_>>();
579        assert!(name[0].is_lowercase(), "Field name must not be capitalized");
580        let res_lower = Ident::new(&name.iter().collect::<String>(), Span::call_site().into());
581        name[0] = name[0].to_ascii_uppercase();
582        let res_upper = Ident::new(&name.iter().collect::<String>(), Span::call_site().into());
583        (res_lower, res_upper)
584    };
585
586    let fields = match &inner.fields {
587        Fields::Named(named) => named.named.iter().collect(),
588        Fields::Unnamed(_) => {
589            return Err(syn::Error::new(
590                name.span(),
591                "Only named fields are supported",
592            ))
593        }
594        Fields::Unit => vec![],
595    };
596
597    let source_field = fields
598        .iter()
599        .filter(|f| f.attrs.iter().any(|attr| attr.path().is_ident("config")))
600        .exactly_one()
601        .map_err(|_| {
602            syn::Error::new(
603                name.span(),
604                "Exactly one field must have the #[config] attribute",
605            )
606        })?;
607    let (source_name, source_name_upper) =
608        gen_name_with_uppercase_idents(source_field.ident.as_ref().unwrap());
609
610    let extensions = fields
611        .iter()
612        .filter(|f| f.attrs.iter().any(|attr| attr.path().is_ident("extension")))
613        .cloned()
614        .collect::<Vec<_>>();
615
616    let mut executor_enum_fields = Vec::new();
617    let mut create_executors = Vec::new();
618    let mut create_airs = Vec::new();
619    let mut execution_where_predicates: Vec<syn::WherePredicate> = Vec::new();
620    let mut circuit_where_predicates: Vec<syn::WherePredicate> = Vec::new();
621
622    let source_field_ty = source_field.ty.clone();
623
624    for e in extensions.iter() {
625        let (ext_field_name, ext_name_upper) =
626            gen_name_with_uppercase_idents(e.ident.as_ref().expect("field must be named"));
627        let executor_type = parse_executor_type(e, false)?;
628        executor_enum_fields.push(quote! {
629            #[any_enum]
630            #ext_name_upper(#executor_type),
631        });
632        create_executors.push(quote! {
633            let inventory: ::openvm_circuit::arch::ExecutorInventory<Self::Executor> = inventory.extend::<F, _, _>(&self.#ext_field_name)?;
634        });
635        let extension_ty = e.ty.clone();
636        execution_where_predicates.push(parse_quote! {
637            #extension_ty: ::openvm_circuit::arch::VmExecutionExtension<F, Executor = #executor_type>
638        });
639        create_airs.push(quote! {
640            inventory.start_new_extension();
641            ::openvm_circuit::arch::VmCircuitExtension::extend_circuit(&self.#ext_field_name, &mut inventory)?;
642        });
643        circuit_where_predicates.push(parse_quote! {
644            #extension_ty: ::openvm_circuit::arch::VmCircuitExtension<SC>
645        });
646    }
647
648    // The config type always needs <F> due to SystemExecutor
649    let source_executor_type = parse_executor_type(source_field, true)?;
650    execution_where_predicates.push(parse_quote! {
651        #source_field_ty: ::openvm_circuit::arch::VmExecutionConfig<F, Executor = #source_executor_type>
652    });
653    circuit_where_predicates.push(parse_quote! {
654        #source_field_ty: ::openvm_circuit::arch::VmCircuitConfig<SC>
655    });
656    let execution_where_clause = quote! { where #(#execution_where_predicates),* };
657    let circuit_where_clause = quote! { where #(#circuit_where_predicates),* };
658
659    let executor_type = Ident::new(&format!("{}Executor", name), name.span());
660
661    let token_stream = TokenStream::from(quote! {
662        #[derive(
663            Clone,
664            ::derive_more::derive::From,
665            ::openvm_circuit::derive::AnyEnum,
666            ::openvm_circuit::derive::Executor,
667            ::openvm_circuit::derive::MeteredExecutor,
668            ::openvm_circuit::derive::PreflightExecutor,
669        )]
670        pub enum #executor_type<F: openvm_stark_backend::p3_field::Field> {
671            #[any_enum]
672            #source_name_upper(#source_executor_type),
673            #(#executor_enum_fields)*
674        }
675
676        impl<F: openvm_stark_backend::p3_field::Field> ::openvm_circuit::arch::VmExecutionConfig<F> for #name #execution_where_clause {
677            type Executor = #executor_type<F>;
678
679            fn create_executors(
680                &self,
681            ) -> Result<::openvm_circuit::arch::ExecutorInventory<Self::Executor>, ::openvm_circuit::arch::ExecutorInventoryError> {
682                let inventory = self.#source_name.create_executors()?.transmute::<Self::Executor>();
683                #(#create_executors)*
684                Ok(inventory)
685            }
686        }
687
688        impl<SC: openvm_stark_backend::config::StarkGenericConfig> ::openvm_circuit::arch::VmCircuitConfig<SC> for #name #circuit_where_clause {
689            fn create_airs(
690                &self,
691            ) -> Result<::openvm_circuit::arch::AirInventory<SC>, ::openvm_circuit::arch::AirInventoryError> {
692                let mut inventory = self.#source_name.create_airs()?;
693                #(#create_airs)*
694                Ok(inventory)
695            }
696        }
697
698        impl AsRef<SystemConfig> for #name {
699            fn as_ref(&self) -> &SystemConfig {
700                self.#source_name.as_ref()
701            }
702        }
703
704        impl AsMut<SystemConfig> for #name {
705            fn as_mut(&mut self) -> &mut SystemConfig {
706                self.#source_name.as_mut()
707            }
708        }
709    });
710    Ok(token_stream)
711}
712
713// Parse the executor name as either
714// `{type_name}Executor` or whatever the attribute `executor = ` specifies
715// Also determines whether the executor type needs generic parameters
716fn parse_executor_type(
717    f: &Field,
718    default_needs_generics: bool,
719) -> syn::Result<proc_macro2::TokenStream> {
720    // TRACKING ISSUE:
721    // We cannot just use <e.ty.to_token_stream() as VmExecutionExtension<F>>::Executor because of this: <https://github.com/rust-lang/rust/issues/85576>
722    let mut executor_type = None;
723    // Do not unwrap the Result until needed
724    let executor_name = syn::parse_str::<Ident>(&format!("{}Executor", f.ty.to_token_stream()));
725
726    if let Some(attr) = f
727        .attrs
728        .iter()
729        .find(|attr| attr.path().is_ident("extension") || attr.path().is_ident("config"))
730    {
731        match attr.meta {
732            Meta::Path(_) => {}
733            Meta::NameValue(_) => {
734                return Err(syn::Error::new(
735                    f.ty.span(),
736                    "Only `#[config]`, `#[extension]`, `#[config(...)]` or `#[extension(...)]` formats are supported",
737                ))
738            }
739            _ => {
740                let nested = attr
741                    .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)?;
742                for meta in nested {
743                    match meta {
744                        Meta::NameValue(nv) => {
745                            if nv.path.is_ident("executor") {
746                                executor_type = match nv.value {
747                                    syn::Expr::Lit(syn::ExprLit {
748                                        lit: syn::Lit::Str(lit_str), ..
749                                    }) => {
750                                        let executor_type: syn::Type = syn::parse_str(&lit_str.value())?;
751                                        Some(quote! { #executor_type })
752                                    },
753                                    syn::Expr::Path(path) => {
754                                        // Handle identifier paths like `executor = MyExecutor`
755                                        Some(path.to_token_stream())
756                                    },
757                                    _ => {
758                                        return Err(syn::Error::new(
759                                            nv.value.span(),
760                                            "executor value must be a string literal or identifier"
761                                        ));
762                                    }
763                                };
764                            } else if nv.path.is_ident("generics") {
765                                // Parse boolean value for generics
766                                let value_str = nv.value.to_token_stream().to_string();
767                                let needs_generics = match value_str.as_str() {
768                                    "true" => true,
769                                    "false" => false,
770                                    _ => return Err(syn::Error::new(
771                                        nv.value.span(),
772                                        "generics attribute must be either true or false"
773                                    ))
774                                };
775                                let executor_name = executor_name.clone()?;
776                                executor_type = Some(if needs_generics {
777                                    quote! { #executor_name<F> }
778                                } else {
779                                    quote! { #executor_name }
780                                });
781                            } else {
782                                return Err(syn::Error::new(nv.span(), "only executor and generics keys are supported"));
783                            }
784                        }
785                        _ => {
786                            return Err(syn::Error::new(meta.span(), "only name = value format is supported"));
787                        }
788                    }
789                }
790            }
791        }
792    }
793    if let Some(executor_type) = executor_type {
794        Ok(executor_type)
795    } else {
796        let executor_name = executor_name?;
797        Ok(if default_needs_generics {
798            quote! { #executor_name<F> }
799        } else {
800            quote! { #executor_name }
801        })
802    }
803}
804
805/// An attribute procedural macro for creating TCO (Tail Call Optimization) handlers.
806///
807/// This macro generates a handler function that wraps an execute implementation
808/// with tail call optimization using the `become` keyword. It extracts the generics
809/// and where clauses from the original function.
810///
811/// # Usage
812///
813/// Place this attribute above a function definition:
814/// ```
815/// #[create_tco_handler]
816/// unsafe fn execute_e1_impl<F: PrimeField32, CTX, const B_IS_IMM: bool>(
817///     pre_compute: &[u8],
818///     state: &mut VmExecState<F, GuestMemory, CTX>,
819/// ) where
820///     CTX: ExecutionCtxTrait,
821/// {
822///     // function body
823/// }
824/// ```
825///
826/// This will generate a TCO handler function with the same generics and where clauses.
827///
828/// # Safety
829///
830/// Do not use this macro if your function wants to terminate execution without error with a
831/// specific error code. The handler generated by this macro assumes that execution should continue
832/// unless the execute_impl returns an error. This is done for performance to skip an exit code
833/// check.
834#[proc_macro_attribute]
835pub fn create_tco_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
836    #[cfg(feature = "tco")]
837    {
838        tco::tco_impl(item)
839    }
840    #[cfg(not(feature = "tco"))]
841    {
842        item
843    }
844}
openvm_circuit_derive/lib.rs

openvm_circuit_derive/
lib.rs
