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    Ident,
12    Item,
13    ItemFn,
14    ItemMod,
15    Meta,
16    PatType,
17    Path,
18    ReturnType,
19    Stmt,
20    Type,
21};
22
23
24#[proc_macro]
25pub fn verify(input: TokenStream) -> TokenStream {
26    let mut args: Vec<syn::Expr> = syn::parse_macro_input!(input with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated).into_iter().collect();
27    let e = args.pop().unwrap();
28    let sig = args.pop().unwrap();
29    let out = quote! {
30        (#sig).assert_valid(stringify!(#e));
31    };
32    out.into()
33}
34
35#[derive(Clone, Debug, Eq, PartialEq)]
36enum RvnAttr {
37    Annotate(String),
38    Assume,
39    Declare,
40    Define,
41    Verify,
42}
43
44fn path_to_one_str(p: &Path) -> String {
45    assert!(p.segments.len() == 1);
46    p.segments.first().unwrap().ident.to_string()
47}
48
49impl RvnAttr {
50    fn try_from_attr(attr: &Attribute) -> Option<RvnAttr> {
51        match &attr.meta {
52            Meta::Path(p) if p.segments.len() == 1 => {
53                match path_to_one_str(p).as_str() {
54                    "assume" => Some(RvnAttr::Assume),
55                    "declare" => Some(RvnAttr::Declare),
56                    "define" => Some(RvnAttr::Define),
57                    "verify" => Some(RvnAttr::Verify),
58                    _ => None,
59                }
60            }
61            Meta::List(l) if l.path.segments.len() == 1 => {
62                match path_to_one_str(&l.path).as_str() {
63                    "annotate" => {
64                        let fun_name: Ident = l.parse_args().unwrap();
65                        Some(RvnAttr::Annotate(fun_name.to_string()))
66                    }
67                    _ => None,
68                }
69            }
70            _ => None,
71        }
72    }
73}
74
75fn pop_rvn_attr(attrs: &mut Vec<Attribute>) -> Option<RvnAttr> {
76    let mut out = None;
77    attrs.retain_mut(|attr| {
78        match RvnAttr::try_from_attr(attr) {
79            Some(ra) => {
80                match &out {
81                    Some(other) => panic!(
82                        "only one ravencheck command should be used per item, but both {:?} and {:?} were used together",
83                        other,
84                        ra,
85                    ),
86                    None => {}
87                }
88
89                out = Some(ra);
90
91                // Drop all ravencheck commands
92                false
93            }
94
95            // Don't drop anything that wasn't a ravencheck command
96            None => true, 
97        }
98    });
99
100    out
101}
102
103#[derive(Clone, Debug, Eq, PartialEq)]
104enum SigItem {
105    Annotation(String, Block),
106    Axiom(Block),
107    Goal(Block),
108    Sort(String),
109    SortAlias(Ident,Type),
110    DFun(Ident, Vec<PatType>, Type, Block),
111    UFun(String, Vec<String>, String),
112}
113
114fn handle_top_level(attrs: &mut Vec<Attribute>, sig_items: &mut Vec<SigItem>) {
115    attrs.retain_mut(|attr| {
116        match &attr.meta {
117            Meta::List(l) if l.path.segments.len() == 1 => {
118                match l.path.segments.first().unwrap().ident.to_string().as_str() {
119                    "declare_types" => {
120                        let parser =
121                            Punctuated
122                            ::<Ident,syn::Token![,]>
123                            ::parse_separated_nonempty;
124                        let types = parser
125                            .parse2(l.tokens.clone())
126                            .expect("the #[declare_types(..)] attribute expects one or more type names as arguments");
127
128                        for t in types.iter() {
129                            sig_items.push(SigItem::Sort(t.to_string()));
130                        }
131
132                        // Don't keep
133                        false
134                    }
135
136                    // Otherwise, do keep
137                    _ => true,
138                }
139            }
140
141            // Otherwise, do keep
142            _ => true,
143        }
144    })
145}
146
147fn handle_items(items: &mut Vec<Item>, sig_items: &mut Vec<SigItem>) {
148    items.retain_mut(|item| {
149        match item {
150            Item::Fn(f) => {
151                match pop_rvn_attr(&mut f.attrs) {
152                    Some(RvnAttr::Annotate(fname)) => {
153                        sig_items.push(
154                            SigItem::Annotation(fname, *(f.block).clone())
155                        );
156                        false
157                    }
158                    Some(RvnAttr::Assume) => {
159                        sig_items.push(
160                            SigItem::Axiom(*(f.block).clone())
161                        );
162                        false
163                    }
164                    Some(RvnAttr::Declare) => {
165                        let name = f.sig.ident.to_string();
166                        let mut arg_types = Vec::new();
167                        for arg in f.sig.inputs.iter() {
168                            match arg {
169                                FnArg::Typed(a) => match *(a.ty.clone()) {
170                                    Type::Path(typepath) => {
171                                        arg_types.push(
172                                            typepath.path.segments.first().unwrap().ident.to_string()
173                                        );
174                                    }
175                                    t => todo!("Handle {:?}", t),
176                                }
177                                FnArg::Receiver(_) => panic!("
178you can't use 'declare' on a method function (one that takes a 'self' input)"
179                                ),
180                            }
181                        }
182                        let out_type = match &f.sig.output {
183                            ReturnType::Default => "()".to_string(),
184                            ReturnType::Type(_, t) => match *(t.clone()) {
185                                Type::Path(typepath) => {
186                                    typepath.path.segments.first().unwrap().ident.to_string()
187                                }
188                                t => todo!("Handle {:?}", t),
189                            }
190                        };
191                        sig_items.push(SigItem::UFun(name, arg_types, out_type));
192
193                        true
194                    }
195                    Some(RvnAttr::Define) => {
196                        let name = f.sig.ident.clone();
197                        let inputs = f.sig.inputs.iter().cloned().map(|arg| {
198                            match arg {
199                                FnArg::Typed(a) => a,
200                                FnArg::Receiver(_) => panic!("
201you can't use 'define' on a method function (one that takes a 'self' input)"
202                                ),
203                            }
204                        }).collect();
205                        let output = match f.sig.output.clone() {
206                            ReturnType::Default => panic!("
207You must give a return type when using 'define'"
208                            ),
209                            ReturnType::Type(_, t) => *t,
210                        };
211                        let body = (*f.block).clone();
212
213                        sig_items.push(
214                            SigItem::DFun(name, inputs, output, body)
215                        );
216                        true
217                    }
218                    Some(RvnAttr::Verify) => {
219                        sig_items.push(
220                            SigItem::Goal(*(f.block).clone())
221                        );
222                        false
223                    }
224                    _ => true,
225                }
226            }
227            Item::Const(i) => {
228                match pop_rvn_attr(&mut i.attrs) {
229                    _ => true,
230                }
231            }
232            Item::Struct(i) => {
233                match pop_rvn_attr(&mut i.attrs) {
234                    Some(RvnAttr::Declare) => {
235                        sig_items.push(SigItem::Sort(i.ident.to_string()));
236                    }
237                    Some(a) => panic!(
238                        "attr {:?} cannot not be used on a struct definition",
239                        a,
240                    ),
241                    None => {}
242                }
243                true
244            }
245            Item::Type(i) => {
246                match pop_rvn_attr(&mut i.attrs) {
247                    Some(RvnAttr::Declare) => {
248                        sig_items.push(SigItem::Sort(i.ident.to_string()));
249                    }
250                    Some(RvnAttr::Define) => {
251                        sig_items.push(SigItem::SortAlias(
252                            i.ident.clone(),
253                            *(i.ty).clone(),
254                        ));
255                    }
256                    Some(a) => panic!(
257                        "attr {:?} cannot not be used on a type alias definition",
258                        a,
259                    ),
260                    None => {}
261                }
262                true
263            }
264            _ => true,
265        }
266    });
267}
268
269#[proc_macro_attribute]
270pub fn check_module(attrs: TokenStream, input: TokenStream) -> TokenStream {
271
272    // The macro needs to name the crate that CheckedSig comes from,
273    // and that will be different depending on the context that
274    // check_module is called in.
275    let cratename: Ident = if attrs.is_empty() {
276        Ident::new("ravencheck", Span::call_site())
277    } else {
278        // The parse_any parser is needed since the crate name could
279        // be a normal ident or a keyword (like 'crate' or 'super').
280        let parser = Ident::parse_any;
281        parser.parse(attrs).expect("parse crate name")
282    };
283
284    let mut m: ItemMod = match syn::parse(input).expect("parse input") {
285        Item::Mod(m) => m,
286        i => panic!(
287            "'check_module' macro should only be applied to a module, but it was applied to {:?}",
288            i,
289        ),
290    };
291
292    let mut sig_items: Vec<SigItem> = Vec::new();
293
294    // Handle commands within the top-level attributes
295    handle_top_level(&mut m.attrs, &mut sig_items);
296
297    // Handle per-item commands within the module
298    match &mut m.content {
299        Some((_,items)) => {
300            handle_items(items, &mut sig_items);
301
302            // Turn sig_items into statements (of type syn::Stmt)
303            let stmts: Vec<Stmt> = sig_items.into_iter().map(|i| {
304                match i {
305                    SigItem::Annotation(f,b) => {
306                        let b_tks = format!("{}", quote! { #b });
307                        syn::parse((quote! {
308                            sig.add_annotation(#f, #b_tks);
309                        }).into()).unwrap()
310                    }
311                    SigItem::Axiom(b) => {
312                        let b_tks = format!("{}", quote! { #b });
313                        syn::parse((quote! {
314                            sig.add_axiom(#b_tks);
315                        }).into()).unwrap()
316                    }
317                    SigItem::Goal(b) => {
318                        let b_tks = format!("{}", quote! { #b });
319                        syn::parse((quote! {
320                            sig.assert_valid(#b_tks);
321                        }).into()).unwrap()
322                    }
323                    SigItem::Sort(s) => {
324                        syn::parse((quote! {
325                            sig.add_sort(#s);
326                        }).into()).unwrap()
327                    }
328                    SigItem::SortAlias(i, ty) => {
329                        let i_tks = format!("{}", i);
330                        let ty_tks = format!("{}", quote! { #ty });
331                        syn::parse((quote! {
332                            sig.add_alias_from_string(#i_tks, #ty_tks);
333                        }).into()).unwrap()
334                    }
335                    SigItem::DFun(name, inputs, _output, body) => {
336                        let name_tk: String = format!("{}", name);
337                        let body_tk: String = format!("{}", quote! {
338                            |#(#inputs),*| #body
339                        });
340                        syn::parse((quote! {
341                            sig.add_fun(#name_tk, #body_tk);
342                        }).into()).unwrap()
343                    }
344                    SigItem::UFun(name, inputs, output) => {
345                        let i_tks_v: Vec<String> = inputs.into_iter().map(|i| {
346                            format!("{}", i)
347                        }).collect();
348                        let tks = if output.as_str() == "bool" {
349                            quote! {
350                                sig.add_relation(#name, [#(#i_tks_v),*]);
351                            }
352                        } else {
353                            let o_tks = format!("{}", output);
354                            quote! {
355                                sig.declare_op(#name, [#(#i_tks_v),*], #o_tks);
356                            }
357                        };
358                        syn::parse(tks.into()).unwrap()
359                    }
360                }
361            }).collect();
362            let test: ItemFn = syn::parse((quote! {
363                #[test]
364                fn check_properties() {
365                    let mut sig = CheckedSig::empty();
366
367                    // Interpolate 'stmts' here
368                    #(#stmts)*
369                }
370            }).into()).unwrap();
371
372            let test_s = Item::Fn(test);
373
374            let test_mod = quote! {
375                #[cfg(test)]
376                mod ravencheck_tests {
377                    use #cratename::CheckedSig;
378
379                    #test_s
380                }
381            };
382
383            items.push(syn::parse(test_mod.into()).expect("parse test mod"));
384        }
385        None => {}
386    }
387    let out = quote! {
388        #m
389    };
390    out.into()
391}