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    ReturnType,
17    Stmt,
18    Type,
19};
20
21
22#[proc_macro]
23pub fn verify(input: TokenStream) -> TokenStream {
24    // let args: Punctuated<syn::Expr,syn::token::Comma> = syn::parse(input).unwrap();
25    let mut args: Vec<syn::Expr> = syn::parse_macro_input!(input with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated).into_iter().collect();
26    let e = args.pop().unwrap();
27    let sig = args.pop().unwrap();
28    let out = quote! {
29        (#sig).assert_valid(stringify!(#e));
30    };
31    // println!("{}", out);
32    out.into()
33}
34
35#[derive(Clone, Debug, Eq, PartialEq)]
36enum RvnAttr {
37    Annotate(String),
38    Assume,
39    Declare,
40    Verify,
41}
42
43fn pop_rvn_attr(attrs: &mut Vec<Attribute>) -> Option<RvnAttr> {
44    let mut out = None;
45    attrs.retain_mut(|attr| {
46        match &attr.meta {
47            Meta::Path(p) if p.segments.len() == 1 => {
48                match p.segments.first().unwrap().ident.to_string().as_str() {
49                    "assume" => {
50                        assert!(
51                            out == None,
52                            "only one attr should be used per item, but used both {:?} and another",
53                            out.clone().unwrap(),
54                        );
55                        out = Some(RvnAttr::Assume);
56                        false
57                    }
58                    "declare" => {
59                        assert!(
60                            out == None,
61                            "only one attr should be used per item, but used both {:?} and another",
62                            out.clone().unwrap(),
63                        );
64                        out = Some(RvnAttr::Declare);
65                        false
66                    }
67                    "verify" => {
68                        assert!(
69                            out == None,
70                            "only one attr should be used per item, but used both {:?} and another",
71                            out.clone().unwrap(),
72                        );
73                        out = Some(RvnAttr::Verify);
74                        false
75                    }
76                    _ => true,
77                }
78            }
79            Meta::List(l) if l.path.segments.len() == 1 => {
80                match l.path.segments.first().unwrap().ident.to_string().as_str() {
81                    "annotate" => {
82                        out = Some(RvnAttr::Annotate(String::new()));
83                        false
84                    }
85                    _ => true,
86                }
87            }
88            _ => true,
89        }
90    });
91
92    out
93}
94
95#[derive(Clone, Debug, Eq, PartialEq)]
96enum SigItem {
97    // Annotation(String, String),
98    Axiom(Block),
99    Goal(Block),
100    Sort(String),
101    UFun(String, Vec<String>, String),
102}
103
104fn handle_items(items: &mut Vec<Item>, sig_items: &mut Vec<SigItem>) {
105    items.retain_mut(|item| {
106        match item {
107            Item::Fn(f) => {
108                match pop_rvn_attr(&mut f.attrs) {
109                    Some(RvnAttr::Annotate(_)) => false,
110                    Some(RvnAttr::Assume) => {
111                        sig_items.push(
112                            SigItem::Axiom(*(f.block).clone())
113                        );
114                        false
115                    }
116                    Some(RvnAttr::Declare) => {
117                        let name = f.sig.ident.to_string();
118                        let mut arg_types = Vec::new();
119                        for arg in f.sig.inputs.iter() {
120                            match arg {
121                                FnArg::Typed(a) => match *(a.ty.clone()) {
122                                    Type::Path(typepath) => {
123                                        arg_types.push(
124                                            typepath.path.segments.first().unwrap().ident.to_string()
125                                        );
126                                    }
127                                    t => todo!("Handle {:?}", t),
128                                }
129                                // FnArg::Typed(_p) => {
130                                //     arg_types.push(
131                                //         stringify!(_p.ty).to_string()
132                                //     );
133                                // }
134                                FnArg::Receiver(_) => panic!("
135you can't use 'declare' on a method function (one that takes a 'self' input)"
136                                ),
137                            }
138                        }
139                        let out_type = match &f.sig.output {
140                            ReturnType::Default => "()".to_string(),
141                            ReturnType::Type(_, t) => match *(t.clone()) {
142                                Type::Path(typepath) => {
143                                    typepath.path.segments.first().unwrap().ident.to_string()
144                                }
145                                t => todo!("Handle {:?}", t),
146                            }
147                        };
148                        sig_items.push(SigItem::UFun(name, arg_types, out_type));
149
150                        true
151                    }
152                    Some(RvnAttr::Verify) => {
153                        sig_items.push(
154                            SigItem::Goal(*(f.block).clone())
155                        );
156                        false
157                    }
158                    _ => true,
159                }
160            }
161            Item::Const(i) => {
162                match pop_rvn_attr(&mut i.attrs) {
163                    _ => true,
164                }
165            }
166            Item::Struct(i) => {
167                match pop_rvn_attr(&mut i.attrs) {
168                    Some(RvnAttr::Declare) => {
169                        sig_items.push(SigItem::Sort(i.ident.to_string()));
170                    }
171                    Some(a) => panic!(
172                        "attr {:?} cannot not be used on a Struct",
173                        a,
174                    ),
175                    None => {}
176                }
177                true
178            }
179            _ => true,
180        }
181    });
182}
183
184#[proc_macro_attribute]
185pub fn check_module(attrs: TokenStream, input: TokenStream) -> TokenStream {
186
187    // The macro needs to name the crate that CheckedSig comes from,
188    // and that will be different depending on the context that
189    // check_module is called in.
190    let cratename: Ident = if attrs.is_empty() {
191        Ident::new("ravencheck", Span::call_site())
192    } else {
193        // The parse_any parser is needed since the crate name could
194        // be a normal ident or a keyword (like 'crate' or 'super').
195        let parser = Ident::parse_any;
196        parser.parse(attrs).expect("parse crate name")
197    };
198
199    let mut m: ItemMod = match syn::parse(input).expect("parse input") {
200        Item::Mod(m) => m,
201        i => panic!(
202            "'check_module' macro should only be applied to a module, but it was applied to {:?}",
203            i,
204        ),
205    };
206
207    let mut sig_items: Vec<SigItem> = Vec::new();
208
209    match &mut m.content {
210        Some((_,items)) => {
211            handle_items(items, &mut sig_items);
212
213            // Turn sig_items into statements (of type syn::Stmt)
214            let stmts: Vec<Stmt> = sig_items.into_iter().map(|i| {
215                match i {
216                    SigItem::Axiom(b) => {
217                        let b_tks = format!("{}", quote! { #b });
218                        syn::parse((quote! {
219                            sig.add_axiom(#b_tks);
220                        }).into()).unwrap()
221                    }
222                    SigItem::Goal(b) => {
223                        let b_tks = format!("{}", quote! { #b });
224                        syn::parse((quote! {
225                            sig.assert_valid(#b_tks);
226                        }).into()).unwrap()
227                    }
228                    SigItem::Sort(s) => {
229                        syn::parse((quote! {
230                            sig.add_sort(#s);
231                        }).into()).unwrap()
232                    }
233                    SigItem::UFun(name, inputs, output) => {
234                        let i_tks_v: Vec<String> = inputs.into_iter().map(|i| {
235                            format!("{}", i)
236                        }).collect();
237                        let tks = if output.as_str() == "bool" {
238                            quote! {
239                                sig.add_relation(#name, [#(#i_tks_v),*]);
240                            }
241                        } else {
242                            let o_tks = format!("{}", output);
243                            quote! {
244                                sig.declare_op(#name, [#(#i_tks_v),*], #o_tks);
245                            }
246                        };
247                        syn::parse(tks.into()).unwrap()
248                    }
249                }
250            }).collect();
251            let test: ItemFn = syn::parse((quote! {
252                #[test]
253                fn check_properties() {
254                    let mut sig = CheckedSig::empty();
255
256                    // Interpolate 'stmts' here
257                    #(#stmts)*
258                }
259            }).into()).unwrap();
260
261            let test_s = Item::Fn(test);
262
263            let test_mod = quote! {
264                #[cfg(test)]
265                mod ravencheck_tests {
266                    use #cratename::CheckedSig;
267
268                    #test_s
269                }
270            };
271
272            items.push(syn::parse(test_mod.into()).expect("parse test mod"));
273        }
274        None => {}
275    }
276    let out = quote! {
277        #m
278    };
279    out.into()
280}