Refactoring contracts macro logic into separate modules

library/kani_macros/src/sysroot/contracts/bootstrap.rs

@@ -0,0 +1,108 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+//! Special way we handle the first time we encounter a contract attribute on a
+//! function.
+
+use proc_macro2::Span;
+use quote::quote;
+use syn::ItemFn;
+
+use super::{
+    helpers::*,
+    shared::{attach_require_kani_any, identifier_for_generated_function},
+    ContractConditionsData, ContractConditionsHandler,
+};
+
+impl<'a> ContractConditionsHandler<'a> {
+    /// The complex case. We are the first time a contract is handled on this function, so
+    /// we're responsible for
+    ///
+    /// 1. Generating a name for the check function
+    /// 2. Emitting the original, unchanged item and register the check
+    ///    function on it via attribute
+    /// 3. Renaming our item to the new name
+    /// 4. And (minor point) adding #[allow(dead_code)] and
+    ///    #[allow(unused_variables)] to the check function attributes
+    pub fn handle_untouched(&mut self) {
+        // We'll be using this to postfix the generated names for the "check"
+        // and "replace" functions.
+        let item_hash = self.hash.unwrap();
+
+        let original_function_name = self.annotated_fn.sig.ident.clone();
+
+        let check_fn_name =
+            identifier_for_generated_function(&original_function_name, "check", item_hash);
+        let replace_fn_name =
+            identifier_for_generated_function(&original_function_name, "replace", item_hash);
+        let recursion_wrapper_name = identifier_for_generated_function(
+            &original_function_name,
+            "recursion_wrapper",
+            item_hash,
+        );
+
+        // Constructing string literals explicitly here, because `stringify!`
+        // doesn't work. Let's say we have an identifier `check_fn` and we were
+        // to do `quote!(stringify!(check_fn))` to try to have it expand to
+        // `"check_fn"` in the generated code. Then when the next macro parses
+        // this it will *not* see the literal `"check_fn"` as you may expect but
+        // instead the *expression* `stringify!(check_fn)`.
+        let replace_fn_name_str = syn::LitStr::new(&replace_fn_name.to_string(), Span::call_site());
+        let wrapper_fn_name_str =
+            syn::LitStr::new(&self.make_wrapper_name().to_string(), Span::call_site());
+        let recursion_wrapper_name_str =
+            syn::LitStr::new(&recursion_wrapper_name.to_string(), Span::call_site());
+
+        // The order of `attrs` and `kanitool::{checked_with,
+        // is_contract_generated}` is important here, because macros are
+        // expanded outside in. This way other contract annotations in `attrs`
+        // sees those attributes and can use them to determine
+        // `function_state`.
+        //
+        // The same care is taken when we emit check and replace functions.
+        // emit the check function.
+        let is_impl_fn = is_probably_impl_fn(&self.annotated_fn);
+        let ItemFn { attrs, vis, sig, block } = &self.annotated_fn;
+        self.output.extend(quote!(
+            #(#attrs)*
+            #[kanitool::checked_with = #recursion_wrapper_name_str]
+            #[kanitool::replaced_with = #replace_fn_name_str]
+            #[kanitool::inner_check = #wrapper_fn_name_str]
+            #vis #sig {
+                #block
+            }
+        ));
+
+        let mut wrapper_sig = sig.clone();
+        attach_require_kani_any(&mut wrapper_sig);
+        wrapper_sig.ident = recursion_wrapper_name;
+
+        let args = pats_to_idents(&mut wrapper_sig.inputs).collect::<Vec<_>>();
+        let also_args = args.iter();
+        let (call_check, call_replace) = if is_impl_fn {
+            (quote!(Self::#check_fn_name), quote!(Self::#replace_fn_name))
+        } else {
+            (quote!(#check_fn_name), quote!(#replace_fn_name))
+        };
+
+        self.output.extend(quote!(
+            #[allow(dead_code, unused_variables)]
+            #[kanitool::is_contract_generated(recursion_wrapper)]
+            #wrapper_sig {
+                static mut REENTRY: bool = false;
+                if unsafe { REENTRY } {
+                    #call_replace(#(#args),*)
+                } else {
+                    unsafe { REENTRY = true };
+                    let result = #call_check(#(#also_args),*);
+                    unsafe { REENTRY = false };
+                    result
+                }
+            }
+        ));
+
+        self.emit_check_function(Some(check_fn_name));
+        self.emit_replace_function(Some(replace_fn_name));
+        self.emit_augmented_modifies_wrapper();
+    }
+}

library/kani_macros/src/sysroot/contracts/check.rs

@@ -0,0 +1,253 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+//! Logic used for generating the code that checks a contract.
+
+use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
+use quote::quote;
+use syn::{Expr, FnArg, ItemFn, Token};
+
+use super::{
+    helpers::*,
+    shared::{make_unsafe_argument_copies, try_as_result_assign_mut},
+    ContractConditionsData, ContractConditionsHandler,
+};
+
+impl<'a> ContractConditionsHandler<'a> {
+    /// Create the body of a check function.
+    ///
+    /// Wraps the conditions from this attribute around `self.body`.
+    ///
+    /// Mutable because a `modifies` clause may need to extend the inner call to
+    /// the wrapper with new arguments.
+    pub fn make_check_body(&mut self) -> TokenStream2 {
+        let mut inner = self.ensure_bootstrapped_check_body();
+        let Self { attr_copy, .. } = self;
+
+        match &self.condition_type {
+            ContractConditionsData::Requires { attr } => {
+                quote!(
+                    kani::assume(#attr);
+                    #(#inner)*
+                )
+            }
+            ContractConditionsData::Ensures { argument_names, attr } => {
+                let (arg_copies, copy_clean) = make_unsafe_argument_copies(&argument_names);
+
+                // The code that enforces the postconditions and cleans up the shallow
+                // argument copies (with `mem::forget`).
+                let exec_postconditions = quote!(
+                    kani::assert(#attr, stringify!(#attr_copy));
+                    #copy_clean
+                );
+
+                assert!(matches!(
+                    inner.pop(),
+                    Some(syn::Stmt::Expr(syn::Expr::Path(pexpr), None))
+                        if pexpr.path.get_ident().map_or(false, |id| id == "result")
+                ));
+
+                quote!(
+                    #arg_copies
+                    #(#inner)*
+                    #exec_postconditions
+                    result
+                )
+            }
+            ContractConditionsData::Modifies { attr } => {
+                let wrapper_name = self.make_wrapper_name().to_string();
+
+                let wrapper_args = if let Some(wrapper_call_args) =
+                    inner.iter_mut().find_map(|stmt| try_as_wrapper_call_args(stmt, &wrapper_name))
+                {
+                    let wrapper_args = make_wrapper_args(wrapper_call_args.len(), attr.len());
+                    wrapper_call_args
+                        .extend(wrapper_args.clone().map(|a| Expr::Verbatim(quote!(#a))));
+                    wrapper_args
+                } else {
+                    unreachable!(
+                        "Invariant broken, check function did not contain a call to the wrapper function"
+                    )
+                };
+
+                quote!(
+                    #(let #wrapper_args = unsafe { kani::internal::Pointer::decouple_lifetime(&#attr) };)*
+                    #(#inner)*
+                )
+            }
+        }
+    }
+
+    /// Get the sequence of statements of the previous check body or create the default one.
+    fn ensure_bootstrapped_check_body(&self) -> Vec<syn::Stmt> {
+        let wrapper_name = self.make_wrapper_name();
+        let return_type = return_type_to_type(&self.annotated_fn.sig.output);
+        if self.is_first_emit() {
+            let args = exprs_for_args(&self.annotated_fn.sig.inputs);
+            let wrapper_call = if is_probably_impl_fn(self.annotated_fn) {
+                quote!(Self::#wrapper_name)
+            } else {
+                quote!(#wrapper_name)
+            };
+            syn::parse_quote!(
+                let result : #return_type = #wrapper_call(#(#args),*);
+                result
+            )
+        } else {
+            self.annotated_fn.block.stmts.clone()
+        }
+    }
+
+    /// Emit the check function into the output stream.
+    ///
+    /// See [`Self::make_check_body`] for the most interesting parts of this
+    /// function.
+    pub fn emit_check_function(&mut self, override_function_dent: Option<Ident>) {
+        self.emit_common_header();
+
+        if self.function_state.emit_tag_attr() {
+            // If it's the first time we also emit this marker. Again, order is
+            // important so this happens as the last emitted attribute.
+            self.output.extend(quote!(#[kanitool::is_contract_generated(check)]));
+        }
+        let body = self.make_check_body();
+        let mut sig = self.annotated_fn.sig.clone();
+        if let Some(ident) = override_function_dent {
+            sig.ident = ident;
+        }
+        self.output.extend(quote!(
+            #sig {
+                #body
+            }
+        ))
+    }
+
+    /// Emit a modifies wrapper, possibly augmenting a prior, existing one.
+    ///
+    /// We only augment if this clause is a `modifies` clause. In that case we
+    /// expand its signature with one new argument of type `&impl Arbitrary` for
+    /// each expression in the clause.
+    pub fn emit_augmented_modifies_wrapper(&mut self) {
+        if let ContractConditionsData::Modifies { attr } = &self.condition_type {
+            let wrapper_args = make_wrapper_args(self.annotated_fn.sig.inputs.len(), attr.len());
+            let sig = &mut self.annotated_fn.sig;
+            for arg in wrapper_args.clone() {
+                let lifetime = syn::Lifetime { apostrophe: Span::call_site(), ident: arg.clone() };
+                sig.inputs.push(FnArg::Typed(syn::PatType {
+                    attrs: vec![],
+                    colon_token: Token![:](Span::call_site()),
+                    pat: Box::new(syn::Pat::Verbatim(quote!(#arg))),
+                    ty: Box::new(syn::Type::Verbatim(quote!(&#lifetime impl kani::Arbitrary))),
+                }));
+                sig.generics.params.push(syn::GenericParam::Lifetime(syn::LifetimeParam {
+                    lifetime,
+                    colon_token: None,
+                    bounds: Default::default(),
+                    attrs: vec![],
+                }));
+            }
+            self.output.extend(quote!(#[kanitool::modifies(#(#wrapper_args),*)]))
+        }
+        self.emit_common_header();
+
+        if self.function_state.emit_tag_attr() {
+            // If it's the first time we also emit this marker. Again, order is
+            // important so this happens as the last emitted attribute.
+            self.output.extend(quote!(#[kanitool::is_contract_generated(wrapper)]));
+        }
+
+        let name = self.make_wrapper_name();
+        let ItemFn { vis, sig, block, .. } = self.annotated_fn;
+
+        let mut sig = sig.clone();
+        sig.ident = name;
+        self.output.extend(quote!(
+            #vis #sig #block
+        ));
+    }
+}
+
+/// Try to interpret this statement as `let result : <...> = <wrapper_fn_name>(args ...);` and
+/// return a mutable reference to the parameter list.
+fn try_as_wrapper_call_args<'a>(
+    stmt: &'a mut syn::Stmt,
+    wrapper_fn_name: &str,
+) -> Option<&'a mut syn::punctuated::Punctuated<syn::Expr, syn::token::Comma>> {
+    let syn::LocalInit { diverge: None, expr: init_expr, .. } = try_as_result_assign_mut(stmt)?
+    else {
+        return None;
+    };
+
+    match init_expr.as_mut() {
+        Expr::Call(syn::ExprCall { func: box_func, args, .. }) => match box_func.as_ref() {
+            syn::Expr::Path(syn::ExprPath { qself: None, path, .. })
+                if path.get_ident().map_or(false, |id| id == wrapper_fn_name) =>
+            {
+                Some(args)
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+/// Make `num` [`Ident`]s with the names `_wrapper_arg_{i}` with `i` starting at `low` and
+/// increasing by one each time.
+fn make_wrapper_args(low: usize, num: usize) -> impl Iterator<Item = syn::Ident> + Clone {
+    (low..).map(|i| Ident::new(&format!("_wrapper_arg_{i}"), Span::mixed_site())).take(num)
+}
+
+#[cfg(test)]
+mod test {
+    macro_rules! detect_impl_fn {
+        ($expect_pass:expr, $($tt:tt)*) => {{
+            let syntax = stringify!($($tt)*);
+            let ast = syn::parse_str(syntax).unwrap();
+            assert!($expect_pass == super::is_probably_impl_fn(&ast),
+                "Incorrect detection.\nExpected is_impl_fun: {}\nInput Expr; {}\nParsed: {:?}",
+                $expect_pass,
+                syntax,
+                ast
+            );
+        }}
+    }
+
+    #[test]
+    fn detect_impl_fn_by_receiver() {
+        detect_impl_fn!(true, fn self_by_ref(&self, u: usize) -> bool {});
+
+        detect_impl_fn!(true, fn self_by_self(self, u: usize) -> bool {});
+    }
+
+    #[test]
+    fn detect_impl_fn_by_self_ty() {
+        detect_impl_fn!(true, fn self_by_construct(u: usize) -> Self {});
+        detect_impl_fn!(true, fn self_by_wrapped_construct(u: usize) -> Arc<Self> {});
+
+        detect_impl_fn!(true, fn self_by_other_arg(u: usize, slf: Self) {});
+
+        detect_impl_fn!(true, fn self_by_other_wrapped_arg(u: usize, slf: Vec<Self>) {})
+    }
+
+    #[test]
+    fn detect_impl_fn_by_qself() {
+        detect_impl_fn!(
+            true,
+            fn self_by_mention(u: usize) {
+                Self::other(u)
+            }
+        );
+    }
+
+    #[test]
+    fn detect_no_impl_fn() {
+        detect_impl_fn!(
+            false,
+            fn self_by_mention(u: usize) {
+                let self_name = 18;
+                let self_lit = "self";
+                let self_lit = "Self";
+            }
+        );
+    }
+}