mdebray/tvl-depot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
This commit is contained in:
Vincent Ambo 2023-01-16 02:07:24 +03:00 committed by tazjin
parent f12f938166
commit 7d0456fa0e
1 changed files with 98 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

124
tvix/eval/builtin-macros/src/lib.rs
View file

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