snix_eval_builtin_macros/lib.rs
1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro2::Span;
5use quote::{ToTokens, quote, quote_spanned};
6use syn::parse::Parse;
7use syn::spanned::Spanned;
8use syn::{
9 Attribute, FnArg, Ident, Item, ItemMod, LitStr, Meta, Pat, PatIdent, PatType, Token, Type,
10 parse_macro_input, parse_quote, parse_quote_spanned, parse2,
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!(
192 "builtin captures a `state` argument, but no state type was defined"
193 );
194 }
195
196 captures_state = true;
197 }
198 }
199 }
200
201 let mut rewritten_args = std::mem::take(&mut f.sig.inputs)
202 .into_iter()
203 .collect::<Vec<_>>();
204
205 // Split out the value arguments from the static arguments.
206 let split_idx = if captures_state { 2 } else { 1 };
207 let value_args = rewritten_args.split_off(split_idx);
208
209 let builtin_arguments = value_args
210 .into_iter()
211 .map(|arg| {
212 let span = arg.span();
213 let mut strict = true;
214 let mut catch = false;
215 let (name, ty) = match arg {
216 FnArg::Receiver(_) => {
217 return Err(quote_spanned!(span => {
218 compile_error!("unexpected receiver argument in builtin")
219 }))
220 }
221 FnArg::Typed(PatType {
222 mut attrs, pat, ty, ..
223 }) => {
224 attrs.retain(|attr| {
225 attr.path.get_ident().into_iter().any(|id| {
226 if id == "lazy" {
227 strict = false;
228 false
229 } else if id == "catch" {
230 catch = true;
231 false
232 } else {
233 true
234 }
235 })
236 });
237 match pat.as_ref() {
238 Pat::Ident(PatIdent { ident, .. }) => {
239 (ident.clone(), ty.clone())
240 }
241 _ => panic!("ignored value parameters must be named, e.g. `_x` and not just `_`"),
242 }
243 }
244 };
245
246 if catch && !strict {
247 return Err(quote_spanned!(span => {
248 compile_error!("Cannot mix both lazy and catch on the same argument")
249 }));
250 }
251
252 Ok(BuiltinArgument {
253 strict,
254 catch,
255 span,
256 name,
257 ty,
258 })
259 })
260 .collect::<Result<Vec<BuiltinArgument>, _>>();
261
262 let builtin_arguments = match builtin_arguments {
263 Err(err) => return err.into(),
264
265 // reverse argument order, as they are popped from the stack
266 // slice in opposite order
267 Ok(args) => args,
268 };
269
270 // Rewrite the argument to the actual function to take a
271 // `Vec<Value>`, which is then destructured into the
272 // user-defined values in the function header.
273 let sig_span = f.sig.span();
274 rewritten_args.push(parse_quote_spanned!(sig_span=> mut values: Vec<Value>));
275 f.sig.inputs = rewritten_args.into_iter().collect();
276
277 // Rewrite the body of the function to do said argument forcing.
278 //
279 // This is done by creating a new block for each of the
280 // arguments that evaluates it, and wraps the inner block.
281 for arg in &builtin_arguments {
282 let block = &f.block;
283 let ty = &arg.ty;
284 let ident = &arg.name;
285
286 f.block = Box::new(match arg {
287 BuiltinArgument {
288 strict: true,
289 catch: true,
290 ..
291 } => parse_quote_spanned! {
292 arg.span => {
293 let #ident: #ty = snix_eval::generators::request_force(
294 &co, values.pop().expect("Snix bug: builtin called with incorrect number of arguments")
295 ).await;
296 #block
297 }
298 },
299 BuiltinArgument {
300 strict: true,
301 catch: false,
302 ..
303 } => parse_quote_spanned! {
304 arg.span => {
305 let #ident: #ty = snix_eval::generators::request_force(
306 &co, values.pop().expect("Snix bug: builtin called with incorrect number of arguments")
307 ).await;
308 if #ident.is_catchable() {
309 return Ok(#ident);
310 }
311 #block
312 }
313 },
314 BuiltinArgument {
315 strict: false,
316 catch: _,
317 ..
318 } => parse_quote_spanned! {
319 arg.span => {
320 let #ident: #ty = values.pop().expect("Snix bug: builtin called with incorrect number of arguments");
321 #block
322 }
323 },
324 });
325 }
326
327 let fn_name = f.sig.ident.clone();
328 let arg_count = builtin_arguments.len();
329 let docstring = match extract_docstring(&f.attrs) {
330 Some(docs) => quote!(Some(#docs)),
331 None => quote!(None),
332 };
333
334 if captures_state {
335 builtins.push(quote_spanned! { builtin_attr.span() => {
336 let inner_state = state.clone();
337 snix_eval::Builtin::new(
338 #name,
339 #docstring,
340 #arg_count,
341 move |values| Gen::new(|co| snix_eval::generators::pin_generator(#fn_name(inner_state.clone(), co, values))),
342 )
343 }});
344 } else {
345 builtins.push(quote_spanned! { builtin_attr.span() => {
346 snix_eval::Builtin::new(
347 #name,
348 #docstring,
349 #arg_count,
350 |values| Gen::new(|co| snix_eval::generators::pin_generator(#fn_name(co, values))),
351 )
352 }});
353 }
354 }
355 }
356 }
357
358 if let Some(state_type) = state_type {
359 items.push(parse_quote! {
360 pub fn builtins(state: #state_type) -> Vec<(&'static str, Value)> {
361 vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect()
362 }
363 });
364 } else {
365 items.push(parse_quote! {
366 pub fn builtins() -> Vec<(&'static str, Value)> {
367 vec![#(#builtins),*].into_iter().map(|b| (b.name(), Value::Builtin(b))).collect()
368 }
369 });
370 }
371
372 module.into_token_stream().into()
373}