ravencheck_macros/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote;
4
5use syn::{
6    Attribute,
7    Ident,
8    Item,
9    ItemFn,
10    ItemMod,
11    ItemUse,
12    Meta,
13    Path,
14    PathArguments,
15    PathSegment,
16    punctuated::Punctuated,
17    Stmt,
18    Token,
19    UseTree,
20};
21use syn::ext::IdentExt;
22use syn::parse::Parser;
23
24use std::collections::VecDeque;
25
26#[derive(Clone, Debug, Eq, PartialEq)]
27enum RvnItemAttr {
28    Annotate(String),
29    AnnotateMulti,
30    Assume,
31    AssumeFor(String),
32    Declare,
33    Define,
34    Falsify,
35    ForCall(String),
36    ForValues(String),
37    Import,
38    InstRule(String),
39    Phantom,
40    Recursive,
41    ShouldFail,
42    Total,
43    Verify,
44}
45
46fn path_to_one_str(p: &Path) -> Option<String> {
47    if p.segments.len() == 1 {
48        Some(p.segments.first().unwrap().ident.to_string())
49    } else {
50        None
51    }
52}
53
54impl RvnItemAttr {
55    fn try_from_attr(attr: &Attribute) -> Option<RvnItemAttr> {
56        match &attr.meta {
57            Meta::Path(p) if p.segments.len() == 1 => {
58                match path_to_one_str(p).as_deref() {
59                    Some("annotate") => panic!("#[annotate] needs arguments"),
60                    Some("annotate_multi") => Some(RvnItemAttr::AnnotateMulti),
61                    Some("assume") => Some(RvnItemAttr::Assume),
62                    Some("declare") => Some(RvnItemAttr::Declare),
63                    Some("define") => Some(RvnItemAttr::Define),
64                    Some("define_rec") => panic!("#[define_rec] has been replaced by #[define] followed by #[recursive]"),
65                    Some("falsify") => Some(RvnItemAttr::Falsify),
66                    Some("import") => Some(RvnItemAttr::Import),
67                    Some("phantom") => Some(RvnItemAttr::Phantom),
68                    Some("recursive") => Some(RvnItemAttr::Recursive),
69                    Some("should_fail") => Some(RvnItemAttr::ShouldFail),
70                    Some("total") => Some(RvnItemAttr::Total),
71                    Some("verify") => Some(RvnItemAttr::Verify),
72                    _ => None,
73                }
74            }
75            Meta::List(l) if l.path.segments.len() == 1 => {
76                match path_to_one_str(&l.path).as_deref() {
77                    Some("annotate") =>
78                        Some(RvnItemAttr::Annotate(l.tokens.to_string())),
79                    Some("assume") =>
80                        Some(RvnItemAttr::AssumeFor(l.tokens.to_string())),
81                    Some("assume_for") => {
82                        panic!("#[assume_for(..)] has been replaced by #[assume(..)]")
83                        // let fun_name: Ident = l.parse_args().unwrap();
84                        // Some(RvnAttr::AssumeFor(fun_name.to_string()))
85                    }
86                    Some("for_call") =>
87                        Some(RvnItemAttr::ForCall(l.tokens.to_string())),
88                    Some("for_type") => Some(RvnItemAttr::InstRule(l.tokens.to_string())),
89                    Some("for_values") =>
90                        Some(RvnItemAttr::ForValues(l.tokens.to_string())),
91                    _ => None,
92                }
93            }
94            _ => None,
95        }
96    }
97
98    /// Remove any `RvnAttr`s found from the given `Vec` and return a
99    /// new `Vec` containing them.
100    fn extract_from_attrs(attrs: &mut Vec<Attribute>) -> Vec<RvnItemAttr> {
101        let mut out = Vec::new();
102        attrs.retain_mut(|attr| {
103            match RvnItemAttr::try_from_attr(attr) {
104                Some(a) => {
105                    out.push(a);
106                    false
107                }
108                None => true,
109            }
110        });
111        out
112    }
113
114    fn extract_from_item(item: &mut Item) -> Vec<RvnItemAttr> {
115        match item {
116            Item::Const(i) => Self::extract_from_attrs(&mut i.attrs),
117            Item::Enum(i) => Self::extract_from_attrs(&mut i.attrs),
118            Item::Fn(i) => Self::extract_from_attrs(&mut i.attrs),
119            Item::Impl(_) => Vec::new(),
120            Item::Mod(_) => Vec::new(),
121            Item::Struct(i) => Self::extract_from_attrs(&mut i.attrs),
122            Item::Type(i) => Self::extract_from_attrs(&mut i.attrs),
123            Item::Use(i) => Self::extract_from_attrs(&mut i.attrs),
124            // Assume anything we haven't listed here is something
125            // that we totally ignore.
126            _item => Vec::new(),
127        }
128    }
129}
130
131#[derive(Clone, Debug, Eq, PartialEq)]
132enum RccCommand {
133    Annotate(String, ItemFn),
134    AnnotateMulti(bool, Vec<String>, Vec<String>, ItemFn),
135    Assume(Vec<String>, ItemFn),
136    AssumeFor(String, ItemFn),
137    /// The first boolean is `true` if this is a phantom
138    /// declaration. The second is `true` if this should get a
139    /// totality axiom.
140    Declare(Item, bool, bool),
141    DeclareType(Ident, usize),
142    /// The first boolean is `true` if this is a phantome definition,
143    /// and the second boolean is `true` if this is a recursive
144    /// definition. The final boolean is `true` if this should get a
145    /// totality axiom (which only makes sense if it's recursive).
146    Define(Item, bool, bool, bool),
147    Import(ItemUse),
148    /// The boolean is `true` if this should be verified, and `false`
149    /// if this should be falsified.
150    Goal(bool, ItemFn),
151}
152
153impl RccCommand {
154    fn mk_declare_define(
155        ras: Vec<RvnItemAttr>,
156        item: Item,
157        define: bool,
158    ) -> (Option<Self>, Option<Item>) {
159        let mut phantom = false;
160        let mut recursive = false;
161        let mut total = false;
162        for a in ras { match a {
163            RvnItemAttr::Phantom => { phantom = true; },
164            RvnItemAttr::Recursive if define => { recursive = true; },
165            RvnItemAttr::Total => { total = true; },
166            a => panic!(
167                "Unexpected {:?} under #[{}]",
168                a,
169                if define { "define"} else { "declare" },
170            ),
171        }}
172        let ret_item = if !phantom {
173            Some(item.clone())
174        } else {
175            None
176        };
177        if define {
178            (Some(Self::Define(item, phantom, recursive, total)), ret_item)
179        } else {
180            (Some(Self::Declare(item, phantom, total)), ret_item)
181        }
182    }
183
184    fn mk_goal(
185        ras: Vec<RvnItemAttr>,
186        item: Item,
187        should_be_valid: bool,
188    ) -> (Option<Self>, Option<Item>) {
189        assert!(
190            ras.len() == 0,
191            "#[verify] and #[falsify] should not have further attributes beneath them",
192        );
193        let item = match item {
194            Item::Fn(i) => i,
195            item => panic!(
196                "should not use #[verify] or #[falsify] on {:?}",
197                item,
198            ),
199        };
200        (Some(Self::Goal(should_be_valid, item)), None)
201    }
202
203    /// Attempt to extract a `RccCommand` from an [`Item`], also
204    /// returning the original (possibly modified) [`Item`] if it
205    /// should remain in the module and be passed along to the Rust
206    /// compiler.
207    fn from_item(mut item: Item) -> (Option<RccCommand>, Option<Item>) {
208        use RvnItemAttr::*;
209
210        let ras = RvnItemAttr::extract_from_item(&mut item);
211
212        // If there are no Ravencheck attrs, return the item
213        // unchanged.
214        if ras.len() == 0 {
215            return (None, Some(item));
216        }
217        let mut ras = VecDeque::from(ras);
218        match ras.pop_front().unwrap() {
219            Annotate(call) => match item {
220                Item::Fn(i) => {
221                    let c = RccCommand::Annotate(call, i);
222                    (Some(c), None)
223                }
224                item => panic!("Can't use #[annotate(..)] on {:?}", item),
225            }
226            AnnotateMulti => match item {
227                Item::Fn(i) => {
228                    let mut call_lines = Vec::new();
229                    let mut value_lines = Vec::new();
230                    let mut should_fail = false;
231                    for a in ras { match a {
232                        RvnItemAttr::ForCall(c) => { call_lines.push(c); },
233                        RvnItemAttr::ForValues(l) => { value_lines.push(l); },
234                        RvnItemAttr::ShouldFail => { should_fail = true; },
235                        a => panic!(
236                            "Unexpected {:?} on '{}'",
237                            a,
238                            i.sig.ident
239                        ),
240                    }}
241                    let c = RccCommand::AnnotateMulti(
242                        should_fail,
243                        value_lines,
244                        call_lines,
245                        i
246                    );
247                    (Some(c), None)
248                }
249                item => panic!("Can't use #[annotate_multi(..)] on {:?}", item),
250            }
251            Assume => match item {
252                Item::Fn(i) => {
253                    let mut rules = Vec::new();
254                    for a in ras { match a {
255                        RvnItemAttr::InstRule(r) => { rules.push(r); },
256                        a => panic!(
257                            "Unexpected {:?} on '{}'",
258                            a,
259                            i.sig.ident
260                        ),
261                    }}
262                    let c = RccCommand::Assume(rules, i);
263                    (Some(c), None)
264                }
265                item => panic!("Can't use #[assume] on {:?}", item),
266            }
267            AssumeFor(call) => match item {
268                Item::Fn(i) => {
269                    let c = RccCommand::AssumeFor(call, i);
270                    (Some(c), None)
271                }
272                item => panic!("Can't use #[assume(..)] on {:?}", item),
273            }
274            Declare => Self::mk_declare_define(Vec::from(ras), item, false),
275            Define => Self::mk_declare_define(Vec::from(ras), item, true),
276            Falsify => Self::mk_goal(Vec::from(ras), item, false),
277            Import => match item {
278                Item::Use(i) => {
279                    let c = RccCommand::Import(i.clone());
280                    (Some(c), Some(Item::Use(i)))
281                }
282                item => panic!("Can't use #[import] on {:?}", item),
283            }
284            InstRule(_) =>
285                panic!("#[for_type(..)] should only appear under #[assume]"),
286            Phantom =>
287                panic!("#[phantom] should only appear under #[declare] or #[define]"),
288            Verify => Self::mk_goal(Vec::from(ras), item, true),
289            a => todo!("other attrs for from_item: {:?}", a),
290        }
291    }
292
293    fn from_toplevel_attr(attr: &Attribute) -> Vec<Self> {
294        let mut out = Vec::new();
295        match &attr.meta {
296            Meta::List(l) if l.path.segments.len() == 1 => {
297                match path_to_one_str(&l.path).as_deref() {
298                    Some("declare_types") => {
299                        let parser =
300                            Punctuated
301                            ::<Path,syn::Token![,]>
302                            ::parse_separated_nonempty;
303                        let types = parser
304                            .parse2(l.tokens.clone())
305                            .expect("the #[declare_types(..)] attribute expects one or more type names as arguments");
306
307                        for mut p in types.into_iter() {
308                            let seg = p.segments.pop().unwrap().into_value();
309                            let arity = match seg.arguments {
310                                PathArguments::None => 0,
311                                PathArguments::AngleBracketed(a) => a.args.len(),
312                                PathArguments::Parenthesized(..) => panic!("declared types should get angle-bracketed arguments <..>, but {} got parenthesized arguments", seg.ident),
313                            };
314
315                            out.push(RccCommand::DeclareType(seg.ident, arity));
316                        }
317                    }
318                    _ => {}
319                }
320            }
321            _ => {}
322        }
323        out
324    }
325
326    fn extract_from_toplevel_attrs(attrs: &mut Vec<Attribute>) -> Vec<Self> {
327        let mut out = Vec::new();
328        attrs.retain_mut(|attr| {
329            let mut cs = Self::from_toplevel_attr(attr);
330            if cs.len() > 0 {
331                out.append(&mut cs);
332                false
333            } else {
334                true
335            }
336        });
337        out
338    }
339
340    fn extract_from_items(items: &mut Vec<Item>) -> Vec<Self> {
341        let mut commands: Vec<Self> = Vec::new();
342        let mut items_out: Vec<Item> = Vec::new();
343        for item in items.clone() {
344            let (c,i) = Self::from_item(item);
345            if let Some(c) = c {
346                commands.push(c);
347            }
348            if let Some(i) = i {
349                items_out.push(i);
350            }
351        }
352        *items = items_out;
353        commands
354    }
355}
356
357#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
358enum GenMode {
359    Check,
360    Export,
361}
362
363fn generate_stmts(commands: &Vec<RccCommand>, mode: GenMode) -> Vec<Stmt> {
364    let mut out = Vec::new();
365    for c in commands {
366        match c {
367            RccCommand::DeclareType(ident, arity) => {
368                let ident_str = quote!{#ident}.to_string();
369                let s: Stmt = syn::parse2(quote! {
370                    rcc.reg_toplevel_type(#ident_str, #arity);
371                }).unwrap();
372                out.push(s);
373            }
374            RccCommand::Annotate(call_str, item_fn) => {
375                let item_str = quote!{ #item_fn }.to_string();
376                let s: Stmt = syn::parse2(quote! {
377                    rcc.reg_fn_annotate(#call_str, #item_str).unwrap();
378                }).unwrap();
379                out.push(s);
380            }
381            RccCommand::AnnotateMulti(should_fail, value_strs, call_strs, item_fn) => {
382                let item_str = quote!{ #item_fn }.to_string();
383                let s: Stmt = syn::parse2(quote! {
384                    rcc.reg_fn_annotate_multi(
385                        #should_fail,
386                        [#(#value_strs),*],
387                        [#(#call_strs),*],
388                        #item_str
389                    ).unwrap();
390                }).unwrap();
391                out.push(s);
392            }
393            RccCommand::Assume(rules, item_fn) => {
394                let item_str = quote!{ #item_fn }.to_string();
395                let s: Stmt = syn::parse2(quote! {
396                    rcc.reg_fn_assume([#(#rules),*], #item_str);
397                }).unwrap();
398                out.push(s);
399            }
400            RccCommand::AssumeFor(call_str, item_fn) => {
401                let item_str = quote!{ #item_fn }.to_string();
402                let s: Stmt = syn::parse2(quote! {
403                    rcc.reg_fn_assume_for(#call_str, #item_str);
404                }).unwrap();
405                out.push(s);
406            }
407            // The effect of 'is_phantom' has already been handled at
408            // the attribute extraction stage.
409            RccCommand::Declare(item, _is_phantom, is_total) => {
410                let item_str = quote!{ #item }.to_string();
411                let s: Stmt = syn::parse2(quote! {
412                    rcc.reg_item_declare(#item_str, #is_total);
413                }).unwrap();
414                out.push(s)
415            }
416            RccCommand::Define(item, _is_phantom, is_rec, is_total) => {
417                let item_str = quote!{ #item }.to_string();
418                let s: Stmt = syn::parse2(quote! {
419                    rcc.reg_item_define(#item_str, #is_rec, #is_total);
420                }).unwrap();
421                out.push(s)
422            }
423            RccCommand::Import(item) => {
424                let segs = use_tree_to_path(item.tree.clone());
425                let mut punct = Punctuated::<PathSegment, Token![::]>::new();
426                for s in segs {
427                    punct.push(s);
428                }
429                let path = Path { leading_colon: None, segments: punct };
430                let path_str = quote!{ #path }.to_string();
431                let s: Stmt = syn::parse2(quote! {
432                    if rcc.touch_new_path(#path_str) {
433                        #path::ravencheck_exports::apply_exports(&mut rcc);
434                    }
435                }).unwrap();
436                out.push(s);
437            }
438            RccCommand::Goal(should_be_valid, item_fn) => {
439                match mode {
440                    GenMode::Check => {
441                        let item_fn_str = quote!{ #item_fn }.to_string();
442                        let s: Stmt = syn::parse2(quote! {
443                            rcc.reg_fn_goal(#should_be_valid, #item_fn_str);
444                        }).unwrap();
445                        out.push(s);
446                    }
447                    GenMode::Export => {}
448                }
449            }
450        }
451    }
452    out
453}
454
455fn use_tree_to_path(t: UseTree) -> Vec<PathSegment> {
456    match t {
457        UseTree::Path(p) => {
458            let i = p.ident.clone();
459            let mut rest = use_tree_to_path(*p.tree);
460            rest.insert(0, PathSegment{
461                ident: i,
462                arguments: PathArguments::None,
463            });
464            rest
465        },
466        UseTree::Glob(..) => Vec::new(),
467        UseTree::Group(..) => Vec::new(),
468        UseTree::Name(..) => Vec::new(),
469        t => panic!("cannot #[import] use-tree with {:?} node in it", t),
470    }
471}
472
473#[proc_macro_attribute]
474pub fn export_module(attrs: TokenStream, input: TokenStream) -> TokenStream {
475    process_module(attrs, input, true)
476}
477
478#[proc_macro_attribute]
479pub fn check_module(attrs: TokenStream, input: TokenStream) -> TokenStream {
480    process_module(attrs, input, false)
481}
482
483fn process_module(
484    attrs: TokenStream,
485    input: TokenStream,
486    export: bool,
487) -> TokenStream {
488
489    // The macro needs to name the crate that CheckedSig comes from,
490    // and that will be different depending on the context that
491    // check_module is called in.
492    let cratename: Ident = if attrs.is_empty() {
493        Ident::new("ravencheck", Span::call_site())
494    } else {
495        // The parse_any parser is needed since the crate name could
496        // be a normal ident or a keyword (like 'crate' or 'super').
497        let parser = Ident::parse_any;
498        parser.parse(attrs).expect("parse crate name")
499    };
500
501    let mut m: ItemMod = match syn::parse(input).expect("parse input") {
502        Item::Mod(m) => m,
503        i => panic!(
504            "'check_module' macro should only be applied to a module, but it was applied to {:?}",
505            i,
506        ),
507    };
508
509    // Handle commands within the top-level attributes
510    let mut commands =
511        RccCommand::extract_from_toplevel_attrs(&mut m.attrs);
512    // extract_top_level(&mut m.attrs, &mut rcc_items);
513
514    // Handle per-item commands within the module
515    match &mut m.content {
516        Some((_,items)) => {
517            
518            // extract_items(items, &mut rcc_items);
519            commands.append(
520                &mut RccCommand::extract_from_items(items)
521            );
522
523            let mut test_stmts: Vec<Stmt> =
524                generate_stmts(&commands, GenMode::Check);
525            test_stmts.push(
526                syn::parse2(quote!{
527                    rcc.check_goals();
528                }).unwrap()
529            );
530
531            let test: ItemFn = syn::parse((quote! {
532                #[test]
533                fn check_properties() {
534                    let mut rcc = Rcc::new();
535
536                    // Interpolate 'stmts' here
537                    #(#test_stmts)*
538                }
539            }).into()).unwrap();
540
541            let test_s = Item::Fn(test);
542
543            let test_mod = quote! {
544                #[cfg(test)]
545                mod ravencheck_tests {
546                    use #cratename::Rcc;
547
548                    #test_s
549                }
550            };
551
552            // println!("Here is the test module content:");
553            // println!("{}", test_mod);
554
555            items.push(syn::parse(test_mod.into()).expect("parse test mod"));
556
557            if export {
558                let export_stmts: Vec<Stmt> =
559                    generate_stmts(&commands, GenMode::Export);
560
561                let export_mod = quote! {
562                    pub mod ravencheck_exports {
563                        use #cratename::Rcc;
564
565                        pub fn apply_exports(rcc: &mut Rcc) {
566                            #(#export_stmts)*
567                        }
568                    }
569                };
570
571                // println!("Here is the export module content:");
572                // println!("{}", export_mod);
573                items.push(syn::parse(export_mod.into()).expect("parse export mod"));
574            }
575        }
576        None => {}
577    }
578    let out = quote! {
579        #m
580    };
581    out.into()
582}