Enable contracts in associated functions

[celinval]

Contracts could not be used with associated function unless they used Self. This is because at a proc-macro level, we cannot determine if we are inside a associated function or not, except in cases where Self is used. In cases where the contract generation could not identify an associated function, it would generate incorrect call, triggering a compilation error.

Another problem with our encoding is that users could not annotate trait implementations with contract attributes. This would try to create new functions inside the trait implementation context, which is not allowed.

In order to solve these issues, we decided to wrap contract logic using closures instead of functions. See the discussion in the original issue (#3206) for more details.

The new solution is still split in two:

1. The proc-macro will now expand the code inside the original function to encode all possible scenarios (check, replace, recursive check, and original body).
2. Instead of using stub, we introduced a new transformation pass that chooses which scenario to pick according to the target harness configuration, and cleanup unused logic.

The expanded function will have the following structure:

#[kanitool::recursion_check = "__kani_recursion_check_modify"]
#[kanitool::checked_with = "__kani_check_modify"]
#[kanitool::replaced_with = "__kani_replace_modify"]
#[kanitool::inner_check = "__kani_modifies_modify"]
fn name_fn(ptr: &mut u32) {
    #[kanitool::fn_marker = "kani_register_contract"]
    pub const fn kani_register_contract<T, F: FnOnce() -> T>(f: F) -> T {
        kani::panic("internal error: entered unreachable code: ")
    }
    let kani_contract_mode = kani::internal::mode();
    match kani_contract_mode {
        kani::internal::RECURSION_CHECK => {
            #[kanitool::is_contract_generated(recursion_check)]
            let mut __kani_recursion_check_name_fn = || { /* recursion check body */ };
            kani_register_contract(__kani_recursion_check_modify)
        }
        kani::internal::REPLACE => {
            #[kanitool::is_contract_generated(replace)]
            let mut __kani_replace_name_fn = || { /* replace body */ };
            kani_register_contract(__kani_replace_name_fn)
        }
        kani::internal::SIMPLE_CHECK => {
            #[kanitool::is_contract_generated(check)]
            let mut __kani_check_name_fn = || { /* check body */ };
            kani_register_contract(__kani_check_name_fn)
        }
        _ => { /* original body */ }
    }
}

In runtime, kani::internal::mode() will return kani::internal::ORIGINAL, which runs the original body. The transformation will replace this call by a different assignment in case the function needs to be replaced. The body of the unused closures will be replaced by a UNREACHABLE statement to avoid unnecessary code to be analyzed.

This is still fairly hacky, but hopefully we can cleanup this logic once Rust adds contract support. 🤞

Resolves #3206

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

[feliperodri]

First batch of reviews. So far the test looks good and we have confidence in the pre-existing tests that the refactoring did not break any current behavior. I'm gonna give another pass in the instrumentation.

[JustusAdam]

I wonder if we couldn't vastly simplify the generated code by just inlining the closures into the match and pick a branch based on mode? Since we're already replacing the mode call with a constant, can't we just run like rustc's folding + dead code removal to select the branch for us?

[celinval]

I wonder if we couldn't vastly simplify the generated code by just inlining the closures into the match and pick a branch based on mode? Since we're already replacing the mode call with a constant, can't we just run like rustc's folding + dead code removal to select the branch for us?

Actually, we cannot do that today unless if we want to restrict contract of constant functions to only contain constant code. I remember bumping into this issue while trying to verify the copy intrinsics.

@feliperodri that is also the reason why we need the kani_register_contract function, i.e., we cannot invoke closures from constant functions.

This restriction is documented in the bootstrap.rs file:

// Dummy function used to force the compiler to capture the environment.
// We cannot call closures inside constant functions.
// This function gets replaced by `kani::internal::call_closure`.
#[inline(never)]
#[kanitool::fn_marker = "kani_register_contract"]
const fn kani_register_contract<T, F: FnOnce() -> T>(f: F) -> T {
    unreachable!()
}

[feliperodri]

🚢 it!