diff options
Diffstat (limited to 'tvix/eval/builtin-macros/src/lib.rs')
| -rw-r--r-- | tvix/eval/builtin-macros/src/lib.rs | 124 |
1 files changed, 98 insertions, 26 deletions
diff --git a/tvix/eval/builtin-macros/src/lib.rs b/tvix/eval/builtin-macros/src/lib.rs index 5a54d499124d..ba59cb2e1685 100644 --- a/tvix/eval/builtin-macros/src/lib.rs +++ b/tvix/eval/builtin-macros/src/lib.rs @@ -6,8 +6,8 @@ use quote::{quote, quote_spanned, ToTokens}; use syn::parse::Parse; use syn::spanned::Spanned; use syn::{ - parse2, parse_macro_input, parse_quote, Attribute, FnArg, Ident, Item, ItemMod, LitStr, Pat, - PatIdent, PatType, Token, + parse2, parse_macro_input, parse_quote, Attribute, FnArg, Ident, Item, ItemMod, LitStr, Meta, + Pat, PatIdent, PatType, Token, Type, }; struct BuiltinArgs { @@ -65,11 +65,42 @@ fn extract_docstring(attrs: &[Attribute]) -> Option<String> { }) } +/// Parse arguments to the `builtins` macro itself, such as `#[builtins(state = Rc<State>)]`. +fn parse_module_args(args: TokenStream) -> Option<Type> { + if args.is_empty() { + return None; + } + + let meta: Meta = syn::parse(args).expect("could not parse arguments to `builtins`-attribute"); + let name_value = match meta { + Meta::NameValue(nv) => nv, + _ => panic!("arguments to `builtins`-attribute must be of the form `name = value`"), + }; + + if name_value.path.get_ident().unwrap().to_string() != "state" { + return None; + } + + if let syn::Lit::Str(type_name) = name_value.lit { + let state_type: Type = + syn::parse_str(&type_name.value()).expect("failed to parse builtins state type"); + return Some(state_type); + } + + panic!("state attribute must be a quoted Rust type"); +} + /// Mark the annotated module as a module for defining Nix builtins. /// +/// An optional type definition may be specified as an argument (e.g. `#[builtins(Rc<State>)]`), +/// which will add a parameter to the `builtins` function of that type which is passed to each +/// builtin upon instantiation. Using this, builtins that close over some external state can be +/// written. +/// /// A function `fn builtins() -> Vec<Builtin>` will be defined within the annotated module, /// returning a list of [`tvix_eval::Builtin`] for each function annotated with the `#[builtin]` -/// attribute within the module. +/// attribute within the module. If a `state` type is specified, the `builtins` function will take a +/// value of that type. /// /// Each invocation of the `#[builtin]` annotation within the module should be passed a string /// literal for the name of the builtin. @@ -100,9 +131,12 @@ fn extract_docstring(attrs: &[Attribute]) -> Option<String> { /// } /// ``` #[proc_macro_attribute] -pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream { +pub fn builtins(args: TokenStream, item: TokenStream) -> TokenStream { let mut module = parse_macro_input!(item as ItemMod); + // parse the optional state type, which users might want to pass to builtins + let state_type = parse_module_args(args); + let (_, items) = match &mut module.content { Some(content) => content, None => { @@ -135,11 +169,26 @@ pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream { .into(); } - let builtin_arguments = f - .sig - .inputs - .iter_mut() - .skip(1) + // Determine if this function is taking the state parameter. + let mut args_iter = f.sig.inputs.iter_mut().peekable(); + let mut captures_state = false; + if let Some(FnArg::Typed(PatType { pat, .. })) = args_iter.peek() { + if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() { + if ident.to_string() == "state" { + if state_type.is_none() { + panic!("builtin captures a `state` argument, but no state type was defined"); + } + + captures_state = true; + } + } + } + + // skip state and/or VM args .. + let skip_num = if captures_state { 2 } else { 1 }; + + let builtin_arguments = args_iter + .skip(skip_num) .map(|arg| { let mut strict = true; let name = match arg { @@ -181,7 +230,7 @@ pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream { }; let fn_name = f.sig.ident.clone(); - let num_args = f.sig.inputs.len() - 1; + let num_args = f.sig.inputs.len() - skip_num; let args = (0..num_args) .map(|n| Ident::new(&format!("arg_{n}"), Span::call_site())) .collect::<Vec<_>>(); @@ -193,26 +242,49 @@ pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream { None => quote!(None), }; - builtins.push(quote_spanned! { builtin_attr.span() => { - crate::Builtin::new( - #name, - &[#(#builtin_arguments),*], - #docstring, - |mut args: Vec<crate::Value>, vm: &mut crate::VM| { - #(let #reversed_args = args.pop().unwrap();)* - #fn_name(vm, #(#args),*) - } - ) - }}); + if captures_state { + builtins.push(quote_spanned! { builtin_attr.span() => { + let inner_state = state.clone(); + crate::Builtin::new( + #name, + &[#(#builtin_arguments),*], + #docstring, + move |mut args: Vec<crate::Value>, vm: &mut crate::VM| { + #(let #reversed_args = args.pop().unwrap();)* + #fn_name(inner_state.clone(), vm, #(#args),*) + } + ) + }}); + } else { + builtins.push(quote_spanned! { builtin_attr.span() => { + crate::Builtin::new( + #name, + &[#(#builtin_arguments),*], + #docstring, + |mut args: Vec<crate::Value>, vm: &mut crate::VM| { + #(let #reversed_args = args.pop().unwrap();)* + #fn_name(vm, #(#args),*) + } + ) + }}); + } } } } - items.push(parse_quote! { - pub fn builtins() -> Vec<(&'static str, Value)> { - vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect() - } - }); + if let Some(state_type) = state_type { + items.push(parse_quote! { + pub fn builtins(state: #state_type) -> Vec<(&'static str, Value)> { + vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect() + } + }); + } else { + items.push(parse_quote! { + pub fn builtins() -> Vec<(&'static str, Value)> { + vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect() + } + }); + } module.into_token_stream().into() } |