about summary refs log tree commit diff
path: root/tvix/eval/builtin-macros/src/lib.rs
diff options
context:
space:
mode:
Diffstat (limited to 'tvix/eval/builtin-macros/src/lib.rs')
-rw-r--r--tvix/eval/builtin-macros/src/lib.rs150
1 files changed, 97 insertions, 53 deletions
diff --git a/tvix/eval/builtin-macros/src/lib.rs b/tvix/eval/builtin-macros/src/lib.rs
index ba59cb2e1685..dfd0948c7d9f 100644
--- a/tvix/eval/builtin-macros/src/lib.rs
+++ b/tvix/eval/builtin-macros/src/lib.rs
@@ -6,20 +6,24 @@ 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, Meta,
-    Pat, PatIdent, PatType, Token, Type,
+    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,
+
+    /// Span at which the argument was defined.
+    span: Span,
 }
 
 fn extract_docstring(attrs: &[Attribute]) -> Option<String> {
@@ -97,6 +101,11 @@ fn parse_module_args(args: TokenStream) -> Option<Type> {
 /// 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. If a `state` type is specified, the `builtins` function will take a
@@ -114,10 +123,10 @@ fn parse_module_args(args: TokenStream) -> Option<Type> {
 ///
 /// #[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)
 ///     }
 ///
@@ -125,7 +134,7 @@ fn parse_module_args(args: TokenStream) -> Option<Type> {
 ///     // 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!()
 ///     }
 /// }
@@ -156,7 +165,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(),
                 };
@@ -169,10 +178,11 @@ pub fn builtins(args: TokenStream, item: TokenStream) -> TokenStream {
                     .into();
                 }
 
-                // Determine if this function is taking the state parameter.
-                let mut args_iter = f.sig.inputs.iter_mut().peekable();
+                // 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 Some(FnArg::Typed(PatType { pat, .. })) = args_iter.peek() {
+                if let FnArg::Typed(PatType { pat, .. }) = &f.sig.inputs[0] {
                     if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() {
                         if ident.to_string() == "state" {
                             if state_type.is_none() {
@@ -184,20 +194,28 @@ pub fn builtins(args: TokenStream, item: TokenStream) -> TokenStream {
                     }
                 }
 
-                // skip state and/or VM args ..
-                let skip_num = if captures_state { 2 } else { 1 };
+                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 = args_iter
-                    .skip(skip_num)
+                let builtin_arguments = value_args
+                    .into_iter()
                     .map(|arg| {
+                        let span = arg.span();
                         let mut strict = true;
-                        let name = match arg {
+                        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" {
@@ -209,34 +227,66 @@ pub fn builtins(args: TokenStream, item: TokenStream) -> TokenStream {
                                     })
                                 });
                                 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::BuiltinArgument {
-                                strict: #strict,
-                                name: #name,
-                            }
-                        }))
+                        Ok(BuiltinArgument {
+                            strict,
+                            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() - skip_num;
-                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 {
+                        f.block = Box::new(parse_quote_spanned! {arg.span=> {
+                            let #ident: #ty = 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 = 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),
@@ -247,24 +297,18 @@ pub fn builtins(args: TokenStream, item: TokenStream) -> TokenStream {
                         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),*)
-                            }
+                            #arg_count,
+                            move |values| Gen::new(|co| generators::pin_generator(#fn_name(inner_state.clone(), co, values))),
                         )
                     }});
                 } 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),*)
-                            }
+                            #arg_count,
+                            |values| Gen::new(|co| generators::pin_generator(#fn_name(co, values))),
                         )
                     }});
                 }