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

openvm_circuit_derive/
lib.rs
