snix_eval_builtin_macros/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro2::Span;
5use quote::{quote, quote_spanned, ToTokens};
6use syn::parse::Parse;
7use syn::spanned::Spanned;
8use syn::{
9    parse2, parse_macro_input, parse_quote, parse_quote_spanned, Attribute, FnArg, Ident, Item,
10    ItemMod, LitStr, Meta, Pat, PatIdent, PatType, Token, Type,
11};
12
13/// Description of a single argument passed to a builtin
14struct BuiltinArgument {
15    /// The name of the argument, to be used in docstrings and error messages
16    name: Ident,
17
18    /// Type of the argument.
19    ty: Box<Type>,
20
21    /// Whether the argument should be forced before the underlying builtin
22    /// function is called.
23    strict: bool,
24
25    /// Propagate catchable values as values to the function, rather than short-circuit returning
26    /// them if encountered
27    catch: bool,
28
29    /// Span at which the argument was defined.
30    span: Span,
31}
32
33fn extract_docstring(attrs: &[Attribute]) -> Option<String> {
34    // Rust docstrings are transparently written pre-macro expansion into an attribute that looks
35    // like:
36    //
37    // #[doc = "docstring here"]
38    //
39    // Multi-line docstrings yield multiple attributes in order, which we assemble into a single
40    // string below.
41
42    #[allow(dead_code)]
43    #[derive(Debug)]
44    struct Docstring {
45        eq: Token![=],
46        doc: LitStr,
47    }
48
49    impl Parse for Docstring {
50        fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
51            Ok(Self {
52                eq: input.parse()?,
53                doc: input.parse()?,
54            })
55        }
56    }
57
58    attrs
59        .iter()
60        .filter(|attr| attr.path.get_ident().into_iter().any(|id| id == "doc"))
61        .filter_map(|attr| parse2::<Docstring>(attr.tokens.clone()).ok())
62        .map(|docstring| docstring.doc.value())
63        .reduce(|mut fst, snd| {
64            if snd.is_empty() {
65                // An empty string represents a spacing newline that was added in the
66                // original doc comment.
67                fst.push_str("\n\n");
68            } else {
69                fst.push_str(&snd);
70            }
71
72            fst
73        })
74}
75
76/// Parse arguments to the `builtins` macro itself, such as `#[builtins(state = Rc<State>)]`.
77fn parse_module_args(args: TokenStream) -> Option<Type> {
78    if args.is_empty() {
79        return None;
80    }
81
82    let meta: Meta = syn::parse(args).expect("could not parse arguments to `builtins`-attribute");
83    let name_value = match meta {
84        Meta::NameValue(nv) => nv,
85        _ => panic!("arguments to `builtins`-attribute must be of the form `name = value`"),
86    };
87
88    if *name_value.path.get_ident().unwrap() != "state" {
89        return None;
90    }
91
92    if let syn::Lit::Str(type_name) = name_value.lit {
93        let state_type: Type =
94            syn::parse_str(&type_name.value()).expect("failed to parse builtins state type");
95        return Some(state_type);
96    }
97
98    panic!("state attribute must be a quoted Rust type");
99}
100
101/// Mark the annotated module as a module for defining Nix builtins.
102///
103/// An optional type definition may be specified as an argument (e.g. `#[builtins(Rc<State>)]`),
104/// which will add a parameter to the `builtins` function of that type which is passed to each
105/// builtin upon instantiation. Using this, builtins that close over some external state can be
106/// written.
107///
108/// The type of each function is rewritten to receive a `Vec<Value>`, containing each `Value`
109/// argument that the function receives. The body of functions is accordingly rewritten to "unwrap"
110/// values from this vector and bind them to the correct names, so unless a static error occurs this
111/// transformation is mostly invisible to users of the macro.
112///
113/// A function `fn builtins() -> Vec<Builtin>` will be defined within the annotated module,
114/// returning a list of `snix_eval::Builtin` for each function annotated with the `#[builtin]`
115/// attribute within the module. If a `state` type is specified, the `builtins` function will take a
116/// value of that type.
117///
118/// Each invocation of the `#[builtin]` annotation within the module should be passed a string
119/// literal for the name of the builtin.
120///
121/// # Examples
122/// ```ignore
123/// # use snix_eval;
124/// # use snix_eval_builtin_macros::builtins;
125///
126/// #[builtins]
127/// mod builtins {
128///     use snix_eval::{GenCo, ErrorKind, Value};
129///
130///     #[builtin("identity")]
131///     pub async fn builtin_identity(co: GenCo, x: Value) -> Result<Value, ErrorKind> {
132///         Ok(x)
133///     }
134///
135///     // Builtins can request their argument not be forced before being called by annotating the
136///     // argument with the `#[lazy]` attribute
137///
138///     #[builtin("tryEval")]
139///     pub async fn builtin_try_eval(co: GenCo, #[lazy] x: Value) -> Result<Value, ErrorKind> {
140///         todo!()
141///     }
142/// }
143/// ```
144#[proc_macro_attribute]
145pub fn builtins(args: TokenStream, item: TokenStream) -> TokenStream {
146    let mut module = parse_macro_input!(item as ItemMod);
147
148    // parse the optional state type, which users might want to pass to builtins
149    let state_type = parse_module_args(args);
150
151    let (_, items) = match &mut module.content {
152        Some(content) => content,
153        None => {
154            return (quote_spanned!(module.span() =>
155                compile_error!("Builtin modules must be defined in-line")
156            ))
157            .into();
158        }
159    };
160
161    let mut builtins = vec![];
162    for item in items.iter_mut() {
163        if let Item::Fn(f) = item {
164            if let Some(builtin_attr_pos) = f
165                .attrs
166                .iter()
167                .position(|attr| attr.path.get_ident().iter().any(|id| *id == "builtin"))
168            {
169                let builtin_attr = f.attrs.remove(builtin_attr_pos);
170                let name: LitStr = match builtin_attr.parse_args() {
171                    Ok(args) => args,
172                    Err(err) => return err.into_compile_error().into(),
173                };
174
175                if f.sig.inputs.len() <= 1 {
176                    return (quote_spanned!(
177                        f.sig.inputs.span() =>
178                            compile_error!("Builtin functions must take at least two arguments")
179                    ))
180                    .into();
181                }
182
183                // Inspect the first argument to determine if this function is
184                // taking the state parameter.
185                // TODO(tazjin): add a test in //snix/eval that covers this
186                let mut captures_state = false;
187                if let FnArg::Typed(PatType { pat, .. }) = &f.sig.inputs[0] {
188                    if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() {
189                        if *ident == "state" {
190                            if state_type.is_none() {
191                                panic!("builtin captures a `state` argument, but no state type was defined");
192                            }
193
194                            captures_state = true;
195                        }
196                    }
197                }
198
199                let mut rewritten_args = std::mem::take(&mut f.sig.inputs)
200                    .into_iter()
201                    .collect::<Vec<_>>();
202
203                // Split out the value arguments from the static arguments.
204                let split_idx = if captures_state { 2 } else { 1 };
205                let value_args = rewritten_args.split_off(split_idx);
206
207                let builtin_arguments = value_args
208                    .into_iter()
209                    .map(|arg| {
210                        let span = arg.span();
211                        let mut strict = true;
212                        let mut catch = false;
213                        let (name, ty) = match arg {
214                            FnArg::Receiver(_) => {
215                                return Err(quote_spanned!(span => {
216                                    compile_error!("unexpected receiver argument in builtin")
217                                }))
218                            }
219                            FnArg::Typed(PatType {
220                                mut attrs, pat, ty, ..
221                            }) => {
222                                attrs.retain(|attr| {
223                                    attr.path.get_ident().into_iter().any(|id| {
224                                        if id == "lazy" {
225                                            strict = false;
226                                            false
227                                        } else if id == "catch" {
228                                            catch = true;
229                                            false
230                                        } else {
231                                            true
232                                        }
233                                    })
234                                });
235                                match pat.as_ref() {
236                                    Pat::Ident(PatIdent { ident, .. }) => {
237                                        (ident.clone(), ty.clone())
238                                    }
239                                    _ => panic!("ignored value parameters must be named, e.g. `_x` and not just `_`"),
240                                }
241                            }
242                        };
243
244                        if catch && !strict {
245                            return Err(quote_spanned!(span => {
246                                compile_error!("Cannot mix both lazy and catch on the same argument")
247                            }));
248                        }
249
250                        Ok(BuiltinArgument {
251                            strict,
252                            catch,
253                            span,
254                            name,
255                            ty,
256                        })
257                    })
258                    .collect::<Result<Vec<BuiltinArgument>, _>>();
259
260                let builtin_arguments = match builtin_arguments {
261                    Err(err) => return err.into(),
262
263                    // reverse argument order, as they are popped from the stack
264                    // slice in opposite order
265                    Ok(args) => args,
266                };
267
268                // Rewrite the argument to the actual function to take a
269                // `Vec<Value>`, which is then destructured into the
270                // user-defined values in the function header.
271                let sig_span = f.sig.span();
272                rewritten_args.push(parse_quote_spanned!(sig_span=> mut values: Vec<Value>));
273                f.sig.inputs = rewritten_args.into_iter().collect();
274
275                // Rewrite the body of the function to do said argument forcing.
276                //
277                // This is done by creating a new block for each of the
278                // arguments that evaluates it, and wraps the inner block.
279                for arg in &builtin_arguments {
280                    let block = &f.block;
281                    let ty = &arg.ty;
282                    let ident = &arg.name;
283
284                    f.block = Box::new(match arg {
285                        BuiltinArgument {
286                            strict: true,
287                            catch: true,
288                            ..
289                        } => parse_quote_spanned! {
290                            arg.span => {
291                                let #ident: #ty = snix_eval::generators::request_force(
292                                    &co, values.pop().expect("Snix bug: builtin called with incorrect number of arguments")
293                                ).await;
294                                #block
295                            }
296                        },
297                        BuiltinArgument {
298                            strict: true,
299                            catch: false,
300                            ..
301                        } => parse_quote_spanned! {
302                            arg.span => {
303                                let #ident: #ty = snix_eval::generators::request_force(
304                                    &co, values.pop().expect("Snix bug: builtin called with incorrect number of arguments")
305                                ).await;
306                                if #ident.is_catchable() {
307                                    return Ok(#ident);
308                                }
309                                #block
310                            }
311                        },
312                        BuiltinArgument {
313                            strict: false,
314                            catch: _,
315                            ..
316                        } => parse_quote_spanned! {
317                            arg.span => {
318                                let #ident: #ty = values.pop().expect("Snix bug: builtin called with incorrect number of arguments");
319                                #block
320                            }
321                        },
322                    });
323                }
324
325                let fn_name = f.sig.ident.clone();
326                let arg_count = builtin_arguments.len();
327                let docstring = match extract_docstring(&f.attrs) {
328                    Some(docs) => quote!(Some(#docs)),
329                    None => quote!(None),
330                };
331
332                if captures_state {
333                    builtins.push(quote_spanned! { builtin_attr.span() => {
334                        let inner_state = state.clone();
335                        snix_eval::Builtin::new(
336                            #name,
337                            #docstring,
338                            #arg_count,
339                            move |values| Gen::new(|co| snix_eval::generators::pin_generator(#fn_name(inner_state.clone(), co, values))),
340                        )
341                    }});
342                } else {
343                    builtins.push(quote_spanned! { builtin_attr.span() => {
344                        snix_eval::Builtin::new(
345                            #name,
346                            #docstring,
347                            #arg_count,
348                            |values| Gen::new(|co| snix_eval::generators::pin_generator(#fn_name(co, values))),
349                        )
350                    }});
351                }
352            }
353        }
354    }
355
356    if let Some(state_type) = state_type {
357        items.push(parse_quote! {
358            pub fn builtins(state: #state_type) -> Vec<(&'static str, Value)> {
359                vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect()
360            }
361        });
362    } else {
363        items.push(parse_quote! {
364            pub fn builtins() -> Vec<(&'static str, Value)> {
365                vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect()
366            }
367        });
368    }
369
370    module.into_token_stream().into()
371}