feat(tvix/eval): let builtin macro capture external state r/5708

This adds a feature to the `#[builtins]` macro which lets users
specify an additional state type to (optionally) thread through to
builtins when constructing them.

This makes it possible for builtins-macro users to pass external state
handles (specifically, in our case, known path tracking) into a set of
builtins.

Change-Id: I3ade20d333fc3ba90a80822cdfa5f87a9cfada75
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7840
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

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()
 }