diff options
Diffstat (limited to 'tvix/eval/builtin-macros/src/lib.rs')
| -rw-r--r-- | tvix/eval/builtin-macros/src/lib.rs | 283 |
1 files changed, 218 insertions, 65 deletions
diff --git a/tvix/eval/builtin-macros/src/lib.rs b/tvix/eval/builtin-macros/src/lib.rs index e815749ec6..5cc9807f54 100644 --- a/tvix/eval/builtin-macros/src/lib.rs +++ b/tvix/eval/builtin-macros/src/lib.rs @@ -6,27 +6,38 @@ 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, parse_quote_spanned, Attribute, FnArg, Ident, Item, + ItemMod, LitStr, Meta, Pat, PatIdent, PatType, Token, Type, }; -struct BuiltinArgs { - name: LitStr, -} +/// Description of a single argument passed to a builtin +struct BuiltinArgument { + /// The name of the argument, to be used in docstrings and error messages + name: Ident, -impl Parse for BuiltinArgs { - fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> { - Ok(BuiltinArgs { - name: input.parse()?, - }) - } + /// Type of the argument. + ty: Box<Type>, + + /// Whether the argument should be forced before the underlying builtin + /// function is called. + strict: bool, + + /// Propagate catchable values as values to the function, rather than short-circuit returning + /// them if encountered + catch: bool, + + /// Span at which the argument was defined. + span: Span, } -fn extract_docstring(attrs: &[Attribute]) -> Option<LitStr> { +fn extract_docstring(attrs: &[Attribute]) -> Option<String> { // Rust docstrings are transparently written pre-macro expansion into an attribute that looks // like: // // #[doc = "docstring here"] + // + // Multi-line docstrings yield multiple attributes in order, which we assemble into a single + // string below. #[allow(dead_code)] #[derive(Debug)] @@ -47,32 +58,77 @@ fn extract_docstring(attrs: &[Attribute]) -> Option<LitStr> { attrs .iter() .filter(|attr| attr.path.get_ident().into_iter().any(|id| id == "doc")) - .find_map(|attr| parse2::<Docstring>(attr.tokens.clone()).ok()) - .map(|docstring| docstring.doc) + .filter_map(|attr| parse2::<Docstring>(attr.tokens.clone()).ok()) + .map(|docstring| docstring.doc.value()) + .reduce(|mut fst, snd| { + if snd.is_empty() { + // An empty string represents a spacing newline that was added in the + // original doc comment. + fst.push_str("\n\n"); + } else { + fst.push_str(&snd); + } + + fst + }) +} + +/// 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() != "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. +/// +/// The type of each function is rewritten to receive a `Vec<Value>`, containing each `Value` +/// argument that the function receives. The body of functions is accordingly rewritten to "unwrap" +/// values from this vector and bind them to the correct names, so unless a static error occurs this +/// transformation is mostly invisible to users of the macro. +/// /// 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. /// /// # Examples /// ```ignore +/// # use tvix_eval; /// # use tvix_eval_builtin_macros::builtins; -/// # mod value { -/// # pub use tvix_eval::Builtin; -/// # } /// /// #[builtins] /// mod builtins { -/// use tvix_eval::{ErrorKind, Value, VM}; +/// use tvix_eval::{GenCo, ErrorKind, Value}; /// /// #[builtin("identity")] -/// pub fn builtin_identity(_vm: &mut VM, x: Value) -> Result<Value, ErrorKind> { +/// pub async fn builtin_identity(co: GenCo, x: Value) -> Result<Value, ErrorKind> { /// Ok(x) /// } /// @@ -80,15 +136,18 @@ fn extract_docstring(attrs: &[Attribute]) -> Option<LitStr> { /// // argument with the `#[lazy]` attribute /// /// #[builtin("tryEval")] -/// pub fn builtin_try_eval(vm: &mut VM, #[lazy] x: Value) -> Result<Value, ErrorKind> { +/// pub async fn builtin_try_eval(co: GenCo, #[lazy] x: Value) -> Result<Value, ErrorKind> { /// todo!() /// } /// } /// ``` #[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 => { @@ -108,7 +167,7 @@ pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream { .position(|attr| attr.path.get_ident().iter().any(|id| *id == "builtin")) { let builtin_attr = f.attrs.remove(builtin_attr_pos); - let BuiltinArgs { name } = match builtin_attr.parse_args() { + let name: LitStr = match builtin_attr.parse_args() { Ok(args) => args, Err(err) => return err.into_compile_error().into(), }; @@ -121,84 +180,178 @@ pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream { .into(); } - let builtin_arguments = f - .sig - .inputs - .iter_mut() - .skip(1) + // Inspect the first argument to determine if this function is + // taking the state parameter. + // TODO(tazjin): add a test in //tvix/eval that covers this + let mut captures_state = false; + if let FnArg::Typed(PatType { pat, .. }) = &f.sig.inputs[0] { + if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() { + if *ident == "state" { + if state_type.is_none() { + panic!("builtin captures a `state` argument, but no state type was defined"); + } + + captures_state = true; + } + } + } + + let mut rewritten_args = std::mem::take(&mut f.sig.inputs) + .into_iter() + .collect::<Vec<_>>(); + + // Split out the value arguments from the static arguments. + let split_idx = if captures_state { 2 } else { 1 }; + let value_args = rewritten_args.split_off(split_idx); + + let builtin_arguments = value_args + .into_iter() .map(|arg| { + let span = arg.span(); let mut strict = true; - let name = match arg { + let mut catch = false; + let (name, ty) = match arg { FnArg::Receiver(_) => { - return Err(quote_spanned!(arg.span() => { - compile_error!("Unexpected receiver argument in builtin") + return Err(quote_spanned!(span => { + compile_error!("unexpected receiver argument in builtin") })) } - FnArg::Typed(PatType { attrs, pat, .. }) => { + FnArg::Typed(PatType { + mut attrs, pat, ty, .. + }) => { attrs.retain(|attr| { attr.path.get_ident().into_iter().any(|id| { if id == "lazy" { strict = false; false + } else if id == "catch" { + catch = true; + false } else { true } }) }); match pat.as_ref() { - Pat::Ident(PatIdent { ident, .. }) => ident.to_string(), - _ => "unknown".to_string(), + Pat::Ident(PatIdent { ident, .. }) => { + (ident.clone(), ty.clone()) + } + _ => panic!("ignored value parameters must be named, e.g. `_x` and not just `_`"), } } }; - Ok(quote_spanned!(arg.span() => { - crate::internal::BuiltinArgument { - strict: #strict, - name: #name, - } - })) + if catch && !strict { + return Err(quote_spanned!(span => { + compile_error!("Cannot mix both lazy and catch on the same argument") + })); + } + + Ok(BuiltinArgument { + strict, + catch, + span, + name, + ty, + }) }) - .collect::<Result<Vec<_>, _>>(); + .collect::<Result<Vec<BuiltinArgument>, _>>(); let builtin_arguments = match builtin_arguments { - Ok(args) => args, Err(err) => return err.into(), + + // reverse argument order, as they are popped from the stack + // slice in opposite order + Ok(args) => args, }; - let fn_name = f.sig.ident.clone(); - let num_args = f.sig.inputs.len() - 1; - let args = (0..num_args) - .map(|n| Ident::new(&format!("arg_{n}"), Span::call_site())) - .collect::<Vec<_>>(); - let mut reversed_args = args.clone(); - reversed_args.reverse(); + // Rewrite the argument to the actual function to take a + // `Vec<Value>`, which is then destructured into the + // user-defined values in the function header. + let sig_span = f.sig.span(); + rewritten_args.push(parse_quote_spanned!(sig_span=> mut values: Vec<Value>)); + f.sig.inputs = rewritten_args.into_iter().collect(); + // Rewrite the body of the function to do said argument forcing. + // + // This is done by creating a new block for each of the + // arguments that evaluates it, and wraps the inner block. + for arg in &builtin_arguments { + let block = &f.block; + let ty = &arg.ty; + let ident = &arg.name; + + if arg.strict { + if arg.catch { + f.block = Box::new(parse_quote_spanned! {arg.span=> { + let #ident: #ty = tvix_eval::generators::request_force(&co, values.pop() + .expect("Tvix bug: builtin called with incorrect number of arguments")).await; + #block + }}); + } else { + f.block = Box::new(parse_quote_spanned! {arg.span=> { + let #ident: #ty = tvix_eval::generators::request_force(&co, values.pop() + .expect("Tvix bug: builtin called with incorrect number of arguments")).await; + if #ident.is_catchable() { + return Ok(#ident); + } + #block + }}); + } + } else { + f.block = Box::new(parse_quote_spanned! {arg.span=> { + let #ident: #ty = values.pop() + .expect("Tvix bug: builtin called with incorrect number of arguments"); + + #block + }}) + } + } + + let fn_name = f.sig.ident.clone(); + let arg_count = builtin_arguments.len(); let docstring = match extract_docstring(&f.attrs) { Some(docs) => quote!(Some(#docs)), None => quote!(None), }; - builtins.push(quote_spanned! { builtin_attr.span() => { - crate::internal::Builtin::new( - #name, - &[#(#builtin_arguments),*], - #docstring, - |mut args: Vec<crate::Value>, vm: &mut crate::internal::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(); + tvix_eval::Builtin::new( + #name, + #docstring, + #arg_count, + move |values| Gen::new(|co| tvix_eval::generators::pin_generator(#fn_name(inner_state.clone(), co, values))), + ) + }}); + } else { + builtins.push(quote_spanned! { builtin_attr.span() => { + tvix_eval::Builtin::new( + #name, + #docstring, + #arg_count, + |values| Gen::new(|co| tvix_eval::generators::pin_generator(#fn_name(co, values))), + ) + }}); + } } } } - items.push(parse_quote! { - pub fn builtins() -> Vec<crate::internal::Builtin> { - vec![#(#builtins),*] - } - }); + 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() } |