Refactoring contracts macro logic into separate modules (#2992)

Refactors the `contracts` module into several smaller modules grouped by
task. The previous modules was very long and interleaved a lot of logic
for various parts of the contracts state machine. This refactoring
groups functions by specificity, task correlation and reusability.

All code changes here are exclusively moves with no changes to
functionality.

The split is as follows:

- `mod.rs` contains definitions of the data structures used to handle
contracts, the entry point for contract handling, as well as a
documentation overview.
- `bootstrap.rs` is the logic that initialized the central data
structures, such as the handler and the functions state
- `check.rs` is all the logic required to create the function that
checks a contract.
- `replace.rs` is all the logic required to create a function that
abstracts via contract
- `shared.rs` are Kani/contract-specific helper functions that are used
in more than one of `bootstrap`, `check` or `replace`
- `helpers.rs` are Kani/contract-independent helper functions

I was going to make the `old` PR but I shudder at the though of
appending more logic to the extremely long `contracts.rs` file so I
figured I'd split the logic real quick before adding more.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.

Co-authored-by: Felipe R. Monteiro <rms.felipe@gmail.com>

Author: JustusAdam

library/kani_macros/src/sysroot/contracts.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/bootstrap.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";
+            }
+        );
+    }
+}