ravencheck_macros/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote;
4use syn::parse::Parser;
5use syn::punctuated::Punctuated;
6use syn::ext::IdentExt;
7use syn::{
8    Attribute,
9    Block,
10    FnArg,
11    GenericParam,
12    Ident,
13    Item,
14    ItemFn,
15    ItemMod,
16    Meta,
17    PatType,
18    Path,
19    PathArguments,
20    ReturnType,
21    Stmt,
22    Type,
23};
24
25
26#[proc_macro]
27pub fn verify(input: TokenStream) -> TokenStream {
28    let mut args: Vec<syn::Expr> = syn::parse_macro_input!(input with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated).into_iter().collect();
29    let e = args.pop().unwrap();
30    let sig = args.pop().unwrap();
31    let out = quote! {
32        (#sig).assert_valid(stringify!(#e));
33    };
34    out.into()
35}
36
37#[derive(Clone, Debug, Eq, PartialEq)]
38enum RvnAttr {
39    Annotate(String),
40    Assume(Vec<(Type,Type)>),
41    AssumeFor(String),
42    Declare,
43    Define(bool),
44    Verify,
45}
46
47fn path_to_one_str(p: &Path) -> String {
48    assert!(p.segments.len() == 1);
49    p.segments.first().unwrap().ident.to_string()
50}
51
52impl RvnAttr {
53    fn try_from_attr(attr: &Attribute) -> Option<RvnAttr> {
54        match &attr.meta {
55            Meta::Path(p) if p.segments.len() == 1 => {
56                match path_to_one_str(p).as_str() {
57                    "assume" => Some(RvnAttr::Assume(Vec::new())),
58                    "declare" => Some(RvnAttr::Declare),
59                    "define" => Some(RvnAttr::Define(false)),
60                    "define_rec" => Some(RvnAttr::Define(true)),
61                    "verify" => Some(RvnAttr::Verify),
62                    _ => None,
63                }
64            }
65            Meta::List(l) if l.path.segments.len() == 1 => {
66                match path_to_one_str(&l.path).as_str() {
67                    "annotate" => {
68                        let fun_name: Ident = l.parse_args().unwrap();
69                        Some(RvnAttr::Annotate(fun_name.to_string()))
70                    }
71                    "assume" => {
72                        let rule: Type = l.parse_args().unwrap();
73                        match rule {
74                            Type::Tuple(t) => {
75                                let a = t.elems[0].clone();
76                                let b = t.elems[1].clone();
77                                Some(RvnAttr::Assume(vec![(a,b)]))
78                            }
79                            t => todo!("Can't handle inst rule {:?}", t),
80                        }
81                    }
82                    "assume_for" => {
83                        let fun_name: Ident = l.parse_args().unwrap();
84                        Some(RvnAttr::AssumeFor(fun_name.to_string()))
85                    }
86                    _ => None,
87                }
88            }
89            _ => None,
90        }
91    }
92}
93
94fn pop_rvn_attr(attrs: &mut Vec<Attribute>) -> Option<RvnAttr> {
95    let mut out = None;
96    attrs.retain_mut(|attr| {
97        match RvnAttr::try_from_attr(attr) {
98            Some(ra) => {
99                match &out {
100                    Some(other) => panic!(
101                        "only one ravencheck command should be used per item, but both {:?} and {:?} were used together",
102                        other,
103                        ra,
104                    ),
105                    None => {}
106                }
107
108                out = Some(ra);
109
110                // Drop all ravencheck commands
111                false
112            }
113
114            // Don't drop anything that wasn't a ravencheck command
115            None => true, 
116        }
117    });
118
119    out
120}
121
122#[derive(Clone, Debug, Eq, PartialEq)]
123enum SigItem {
124    Annotation(String, String, Block),
125    Axiom(Block, Vec<Ident>, Vec<(Type,Type)>),
126    FunctionAxiom(String, String, Block),
127    Goal(String, Block),
128    TypeAlias(Ident,Type),
129    // The final bool is true if this is a recursive def
130    DFun(Ident, Vec<Ident>, Vec<PatType>, Type, Block, bool),
131    Type(Ident, usize),
132    UFun(String, Vec<Ident>, Vec<PatType>, Type),
133    UConst(String, Type),
134}
135
136fn handle_top_level(attrs: &mut Vec<Attribute>, sig_items: &mut Vec<SigItem>) {
137    attrs.retain_mut(|attr| {
138        match &attr.meta {
139            Meta::List(l) if l.path.segments.len() == 1 => {
140                match l.path.segments.first().unwrap().ident.to_string().as_str() {
141                    "declare_types" => {
142                        let parser =
143                            Punctuated
144                            ::<Path,syn::Token![,]>
145                            ::parse_separated_nonempty;
146                        let types = parser
147                            .parse2(l.tokens.clone())
148                            .expect("the #[declare_types(..)] attribute expects one or more type names as arguments");
149
150                        for mut p in types.into_iter() {
151                            let seg = p.segments.pop().unwrap().into_value();
152                            let args_len = match seg.arguments {
153                                PathArguments::None => 0,
154                                PathArguments::AngleBracketed(a) => a.args.len(),
155                                PathArguments::Parenthesized(..) => panic!("declared types should get angle-bracketed arguments <..>, but {} got parenthesized arguments", seg.ident),
156                            };
157                            sig_items.push(SigItem::Type(seg.ident, args_len));
158                        }
159
160                        // Don't keep
161                        false
162                    }
163
164                    // Otherwise, do keep
165                    _ => true,
166                }
167            }
168
169            // Otherwise, do keep
170            _ => true,
171        }
172    })
173}
174
175fn handle_items(items: &mut Vec<Item>, sig_items: &mut Vec<SigItem>) {
176    items.retain_mut(|item| {
177        match item {
178            Item::Fn(f) => {
179                match pop_rvn_attr(&mut f.attrs) {
180                    Some(RvnAttr::Annotate(fname)) => {
181                        sig_items.push(
182                            SigItem::Annotation(
183                                f.sig.ident.to_string(),
184                                fname,
185                                *(f.block).clone(),
186                            )
187                        );
188                        false
189                    }
190                    Some(RvnAttr::AssumeFor(fname)) => {
191                        sig_items.push(
192                            SigItem::FunctionAxiom(
193                                f.sig.ident.to_string(),
194                                fname,
195                                *(f.block).clone(),
196                            )
197                        );
198                        false
199                    }
200                    Some(RvnAttr::Assume(rules)) => {
201                        let mut tas = Vec::new();
202                        for g in f.sig.generics.params.iter() {
203                            match g {
204                                GenericParam::Type(tp) =>
205                                    tas.push(tp.ident.clone()),
206                                _ => {}
207                            }
208                        }
209                        sig_items.push(
210                            SigItem::Axiom(
211                                *(f.block).clone(),
212                                tas,
213                                rules,
214                            )
215                        );
216                        false
217                    }
218                    Some(RvnAttr::Declare) => {
219                        let name = f.sig.ident.to_string();
220
221                        let mut targs = Vec::new();
222                        for g in f.sig.generics.params.iter() {
223                            match g {
224                                GenericParam::Type(tp) =>
225                                    targs.push(tp.ident.clone()),
226                                _ => {}
227                            }
228                        }
229
230                        let inputs = f.sig.inputs.iter().cloned().map(|arg| {
231                            match arg {
232                                FnArg::Typed(a) => a,
233                                FnArg::Receiver(_) => panic!("
234you can't use 'declare' on a method function (one that takes a 'self' input)"
235                                ),
236                            }
237                        }).collect();
238
239                        let output = match f.sig.output.clone() {
240                            ReturnType::Default => panic!("
241You must give a return type when using 'declare'"
242                            ),
243                            ReturnType::Type(_, t) => *t,
244                        };
245
246                        sig_items.push(SigItem::UFun(name, targs, inputs, output));
247
248                        true
249                    }
250                    Some(RvnAttr::Define(is_rec)) => {
251                        let name = f.sig.ident.clone();
252
253                        let mut targs = Vec::new();
254                        for g in f.sig.generics.params.iter() {
255                            match g {
256                                GenericParam::Type(tp) =>
257                                    targs.push(tp.ident.clone()),
258                                _ => {}
259                            }
260                        }
261
262                        let inputs = f.sig.inputs.iter().cloned().map(|arg| {
263                            match arg {
264                                FnArg::Typed(a) => a,
265                                FnArg::Receiver(_) => panic!("
266you can't use 'define' on a method function (one that takes a 'self' input)"
267                                ),
268                            }
269                        }).collect();
270                        let output = match f.sig.output.clone() {
271                            ReturnType::Default => panic!("
272You must give a return type when using 'define'"
273                            ),
274                            ReturnType::Type(_, t) => *t,
275                        };
276                        let body = (*f.block).clone();
277
278                        sig_items.push(
279                            SigItem::DFun(name, targs, inputs, output, body, is_rec)
280                        );
281                        true
282                    }
283                    // Some(RvnAttr::DefineRec) => {
284                    //     todo!("Handle #[define_rec] for fn");
285                    // }
286                    Some(RvnAttr::Verify) => {
287                        sig_items.push(
288                            SigItem::Goal(f.sig.ident.to_string(), *(f.block).clone())
289                        );
290                        false
291                    }
292                    None => true,
293                }
294            }
295            Item::Const(i) => {
296                match pop_rvn_attr(&mut i.attrs) {
297                    Some(RvnAttr::Declare) => {
298                        sig_items.push(
299                            SigItem::UConst(
300                                i.ident.to_string(),
301                                *(i.ty).clone()
302                            )
303                        );
304                        true
305                    }
306                    _ => true,
307                }
308            }
309            Item::Struct(i) => {
310                match pop_rvn_attr(&mut i.attrs) {
311                    Some(RvnAttr::Declare) => {
312                        let mut num_types = 0;
313                        for a in i.generics.params.iter() {
314                            match a {
315                                GenericParam::Lifetime(..) =>
316                                    panic!("Lifetime params on declared structs are not supported ({})", i.ident),
317                                GenericParam::Type(..) => {
318                                    num_types = num_types + 1;
319                                }
320                                GenericParam::Const(..) =>
321                                    panic!("Const params on declared structs are not supported ({})", i.ident),
322                            }
323                        }
324                        sig_items.push(SigItem::Type(
325                            i.ident.clone(),
326                            num_types,
327                        ));
328                    }
329                    Some(a) => panic!(
330                        "attr {:?} cannot not be used on a struct definition",
331                        a,
332                    ),
333                    None => {}
334                }
335                true
336            }
337            Item::Type(i) => {
338                match pop_rvn_attr(&mut i.attrs) {
339                    Some(RvnAttr::Declare) => {
340                        sig_items.push(SigItem::Type(i.ident.clone(), 0));
341                    }
342                    Some(RvnAttr::Define(_is_rec)) => {
343                        sig_items.push(SigItem::TypeAlias(
344                            i.ident.clone(),
345                            *(i.ty).clone(),
346                        ));
347                    }
348                    Some(a) => panic!(
349                        "attr {:?} cannot not be used on a type alias definition",
350                        a,
351                    ),
352                    None => {}
353                }
354                true
355            }
356            _ => true,
357        }
358    });
359}
360
361#[proc_macro_attribute]
362pub fn check_module(attrs: TokenStream, input: TokenStream) -> TokenStream {
363
364    // The macro needs to name the crate that CheckedSig comes from,
365    // and that will be different depending on the context that
366    // check_module is called in.
367    let cratename: Ident = if attrs.is_empty() {
368        Ident::new("ravencheck", Span::call_site())
369    } else {
370        // The parse_any parser is needed since the crate name could
371        // be a normal ident or a keyword (like 'crate' or 'super').
372        let parser = Ident::parse_any;
373        parser.parse(attrs).expect("parse crate name")
374    };
375
376    let mut m: ItemMod = match syn::parse(input).expect("parse input") {
377        Item::Mod(m) => m,
378        i => panic!(
379            "'check_module' macro should only be applied to a module, but it was applied to {:?}",
380            i,
381        ),
382    };
383
384    let mut sig_items: Vec<SigItem> = Vec::new();
385
386    // Handle commands within the top-level attributes
387    handle_top_level(&mut m.attrs, &mut sig_items);
388
389    // Handle per-item commands within the module
390    match &mut m.content {
391        Some((_,items)) => {
392            handle_items(items, &mut sig_items);
393
394            // Turn sig_items into statements (of type syn::Stmt)
395            let stmts: Vec<Stmt> = sig_items.into_iter().map(|i| {
396                match i {
397                    SigItem::Annotation(title,f,b) => {
398                        let b_tks = format!("{}", quote! { #b });
399                        syn::parse((quote! {
400                            sig.add_checked_annotation(#title, #f, #b_tks);
401                        }).into()).unwrap()
402                    }
403                    SigItem::FunctionAxiom(_title,f,b) => {
404                        let b_tks = format!("{}", quote! { #b });
405                        syn::parse((quote! {
406                            sig.add_annotation(#f, #b_tks);
407                        }).into()).unwrap()
408                    }
409                    SigItem::Axiom(b,tas,rules) => {
410                        let b_tks = format!("{}", quote! { #b });
411                        let tas: Vec<String> = tas.into_iter().map(|i| {
412                            format!("{}", quote! { #i })
413                        }).collect();
414                        let rules: Vec<(String,String)> = rules.into_iter().map(|(a,b)| {
415                            (format!("{}", quote! { #a }), format!("{}", quote! { #b }))
416                        }).collect();
417                        let rules: Vec<_> = rules.into_iter().map(|(a,b)| {
418                            quote! { (#a, #b) }
419                        }).collect();
420                        syn::parse((quote! {
421                            sig.add_axiom2(#b_tks, [#(#tas),*], [#(#rules),*]);
422                        }).into()).unwrap()
423                    }
424                    SigItem::Goal(title, b) => {
425                        let b_tks = format!("{}", quote! { #b });
426                        syn::parse((quote! {
427                            sig.assert_valid_t(#title, #b_tks);
428                        }).into()).unwrap()
429                    }
430                    SigItem::Type(name, arg_len) => {
431                        let s = name.to_string();
432                        syn::parse((quote! {
433                            sig.add_type_con(#s, #arg_len);
434                        }).into()).unwrap()
435                    }
436                    SigItem::TypeAlias(i, ty) => {
437                        let i_tks = format!("{}", i);
438                        let ty_tks = format!("{}", quote! { #ty });
439                        syn::parse((quote! {
440                            sig.add_alias_from_string(#i_tks, #ty_tks);
441                        }).into()).unwrap()
442                    }
443                    SigItem::DFun(name, targs, inputs, output, body, is_rec) => {
444                        let name_tk: String = format!("{}", name);
445                        let targs: Vec<String> = targs.into_iter().map(|i| {
446                            format!("{}", quote! { #i })
447                        }).collect();
448                        let body_tk: String = format!("{}", quote! {
449                            |#(#inputs),*| #body
450                        });
451                        let inputs: Vec<String> = inputs.into_iter().map(|i| {
452                            format!("{}", quote! { #i })
453                        }).collect();
454                        let output: String = format!("{}", quote! { #output });
455                        if !is_rec {
456                            syn::parse((quote! {
457                                sig.add_fun_tas(#name_tk, [#(#targs),*], #body_tk);
458                            }).into()).unwrap()
459                        } else {
460                            syn::parse((quote! {
461                                sig.define_op_rec(#name_tk, [#(#targs),*], [#(#inputs),*], #output, #body_tk);
462                            }).into()).unwrap()
463                        }
464                    }
465                    SigItem::UFun(name, targs, inputs, output) => {
466                        let name: String = format!("{}", name);
467                        let targs: Vec<String> = targs.into_iter().map(|i| {
468                            format!("{}", quote! { #i })
469                        }).collect();
470                        let inputs: Vec<String> = inputs.into_iter().map(|i| {
471                            format!("{}", quote! { #i })
472                        }).collect();
473                        let output: String = format!("{}", quote! { #output });
474                        let tokens = quote! {
475                            sig.declare_op(#name, [#(#targs),*], [#(#inputs),*], #output);
476                        };
477                        // println!("{}", tokens);
478                        syn::parse2(tokens).unwrap()
479                    }
480                    SigItem::UConst(name, ty) => {
481                        let output: String = format!("{}", quote! { #ty });
482                        let tokens = quote! {
483                            sig.declare_const(#name, #output);
484                        };
485                        syn::parse2(tokens).unwrap()
486                    }
487                }
488            }).collect();
489            let test: ItemFn = syn::parse((quote! {
490                #[test]
491                fn check_properties() {
492                    let mut sig = CheckedSig::empty();
493
494                    // Interpolate 'stmts' here
495                    #(#stmts)*
496                }
497            }).into()).unwrap();
498
499            let test_s = Item::Fn(test);
500
501            let test_mod = quote! {
502                #[cfg(test)]
503                mod ravencheck_tests {
504                    use #cratename::CheckedSig;
505
506                    #test_s
507                }
508            };
509
510            items.push(syn::parse(test_mod.into()).expect("parse test mod"));
511        }
512        None => {}
513    }
514    let out = quote! {
515        #m
516    };
517    out.into()
518}