about summary refs log tree commit diff
path: root/tvix/eval/builtin-macros/src/lib.rs
blob: dfd0948c7d9fc59e91345e373b431940f6fcdbcb (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{quote, quote_spanned, ToTokens};
use syn::parse::Parse;
use syn::spanned::Spanned;
use syn::{
    parse2, parse_macro_input, parse_quote, parse_quote_spanned, Attribute, FnArg, Ident, Item,
    ItemMod, LitStr, Meta, Pat, PatIdent, PatType, Token, Type,
};

/// 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,

    /// 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> {
    // 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)]
    struct Docstring {
        eq: Token![=],
        doc: LitStr,
    }

    impl Parse for Docstring {
        fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
            Ok(Self {
                eq: input.parse()?,
                doc: input.parse()?,
            })
        }
    }

    attrs
        .iter()
        .filter(|attr| attr.path.get_ident().into_iter().any(|id| id == "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().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.
///
/// 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
/// 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_builtin_macros::builtins;
/// # mod value {
/// #     pub use tvix_eval::Builtin;
/// # }
///
/// #[builtins]
/// mod builtins {
///     use tvix_eval::{GenCo, ErrorKind, Value};
///
///     #[builtin("identity")]
///     pub async fn builtin_identity(co: GenCo, x: Value) -> Result<Value, ErrorKind> {
///         Ok(x)
///     }
///
///     // Builtins can request their argument not be forced before being called by annotating the
///     // argument with the `#[lazy]` attribute
///
///     #[builtin("tryEval")]
///     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 {
    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 => {
            return (quote_spanned!(module.span() =>
                compile_error!("Builtin modules must be defined in-line")
            ))
            .into();
        }
    };

    let mut builtins = vec![];
    for item in items.iter_mut() {
        if let Item::Fn(f) = item {
            if let Some(builtin_attr_pos) = f
                .attrs
                .iter()
                .position(|attr| attr.path.get_ident().iter().any(|id| *id == "builtin"))
            {
                let builtin_attr = f.attrs.remove(builtin_attr_pos);
                let name: LitStr = match builtin_attr.parse_args() {
                    Ok(args) => args,
                    Err(err) => return err.into_compile_error().into(),
                };

                if f.sig.inputs.len() <= 1 {
                    return (quote_spanned!(
                        f.sig.inputs.span() =>
                            compile_error!("Builtin functions must take at least two arguments")
                    ))
                    .into();
                }

                // 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.to_string() == "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, ty) = match arg {
                            FnArg::Receiver(_) => {
                                return Err(quote_spanned!(span => {
                                    compile_error!("unexpected receiver argument in builtin")
                                }))
                            }
                            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 {
                                            true
                                        }
                                    })
                                });
                                match pat.as_ref() {
                                    Pat::Ident(PatIdent { ident, .. }) => {
                                        (ident.clone(), ty.clone())
                                    }
                                    _ => panic!("ignored value parameters must be named, e.g. `_x` and not just `_`"),
                                }
                            }
                        };

                        Ok(BuiltinArgument {
                            strict,
                            span,
                            name,
                            ty,
                        })
                    })
                    .collect::<Result<Vec<BuiltinArgument>, _>>();

                let builtin_arguments = match builtin_arguments {
                    Err(err) => return err.into(),

                    // reverse argument order, as they are popped from the stack
                    // slice in opposite order
                    Ok(args) => args,
                };

                // 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),
                };

                if captures_state {
                    builtins.push(quote_spanned! { builtin_attr.span() => {
                        let inner_state = state.clone();
                        crate::Builtin::new(
                            #name,
                            #docstring,
                            #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,
                            #docstring,
                            #arg_count,
                            |values| Gen::new(|co| generators::pin_generator(#fn_name(co, values))),
                        )
                    }});
                }
            }
        }
    }

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