Fix contract for std constant functions

Stable constant functions in the standard library can only invoke stable
constant functions. Instead of adding the wrapper in the Kani library,
expand it as an internal function which inherit the attributes of the
function being annotated.

kani-compiler/src/kani_middle/transform/contracts.rs

@@ -400,7 +400,6 @@ impl FunctionWithContractPass {
     /// 2. Replace `kani_register_contract` by the call to the closure.
     fn set_mode(&self, tcx: TyCtxt, body: Body, mode: ContractMode) -> Body {
         debug!(?mode, "set_mode");
-        let mode_fn = find_fn_def(tcx, "KaniContractMode").unwrap();
         let mut new_body = MutableBody::from(body);
         let (mut mode_call, ret, target) = new_body
             .blocks()
@@ -410,16 +409,16 @@ impl FunctionWithContractPass {
                 if let TerminatorKind::Call { func, target, destination, .. } = &bb.terminator.kind
                 {
                     let (callee, _) = func.ty(new_body.locals()).unwrap().kind().fn_def()?;
-                    (callee == mode_fn).then(|| {
-                        (
+                    let marker = KaniAttributes::for_def_id(tcx, callee.def_id()).fn_marker();
+                    if marker.is_some_and(|s| s.as_str() == "kani_contract_mode") {
+                        return Some((
                             SourceInstruction::Terminator { bb: bb_idx },
                             destination.clone(),
                             target.unwrap(),
-                        )
-                    })
-                } else {
-                    None
+                        ));
+                    }
                 }
+                None
             })
             .unwrap();
 

library/kani/src/internal.rs

@@ -143,9 +143,9 @@ fn run_contract_fn<T, F: FnOnce() -> T>(func: F) -> T {
     func()
 }
 
-/// This is used for documentation's sake of which implementation to keep during contract verification.
+/// This is used by contracts to select which version of the contract to use during codegen.
 #[doc(hidden)]
-type Mode = u8;
+pub type Mode = u8;
 
 /// Keep the original body.
 pub const ORIGINAL: Mode = 0;
@@ -158,19 +158,3 @@ pub const SIMPLE_CHECK: Mode = 2;
 
 /// Stub the body with its contract.
 pub const REPLACE: Mode = 3;
-
-/// This function is only used to help with contract instrumentation.
-///
-/// It should be removed from the end user code during contract transformation.
-/// By default, return the original code (used in concrete playback).
-#[doc(hidden)]
-#[inline(never)]
-#[crate::unstable(
-    feature = "function-contracts",
-    issue = 2652,
-    reason = "experimental support for function contracts"
-)]
-#[rustc_diagnostic_item = "KaniContractMode"]
-pub const fn mode() -> Mode {
-    ORIGINAL
-}

library/kani_core/src/lib.rs

@@ -424,9 +424,9 @@ macro_rules! kani_intrinsics {
                 func()
             }
 
-            /// This is used for documentation's sake of which implementation to keep during contract verification.
+            /// This is used by contracts to select which version of the contract to use during codegen.
             #[doc(hidden)]
-            type Mode = u8;
+            pub type Mode = u8;
 
             /// Keep the original body.
             pub const ORIGINAL: Mode = 0;
@@ -439,17 +439,6 @@ macro_rules! kani_intrinsics {
 
             /// Stub the body with its contract.
             pub const REPLACE: Mode = 3;
-
-            /// This function is only used to help with contract instrumentation.
-            ///
-            /// It should be removed from the end user code during contract transformation.
-            /// By default, return the original code (used in concrete playback).
-            #[doc(hidden)]
-            #[inline(never)]
-            #[rustc_diagnostic_item = "KaniContractMode"]
-            pub const fn mode() -> Mode {
-                ORIGINAL
-            }
         }
     };
 }

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

@@ -48,10 +48,17 @@ impl<'a> ContractConditionsHandler<'a> {
                 // This function gets replaced by `kani::internal::call_closure`.
                 #[inline(never)]
                 #[kanitool::fn_marker = "kani_register_contract"]
-                pub const fn kani_register_contract<T, F: FnOnce() -> T>(f: F) -> T {
+                const fn kani_register_contract<T, F: FnOnce() -> T>(f: F) -> T {
                     unreachable!()
                 }
-                let kani_contract_mode = kani::internal::mode();
+                // Dummy function that we replace to pick the contract mode.
+                // By default, return ORIGINAL
+                #[inline(never)]
+                #[kanitool::fn_marker = "kani_contract_mode"]
+                const fn kani_contract_mode() -> kani::internal::Mode {
+                    kani::internal::ORIGINAL
+                }
+                let kani_contract_mode = kani_contract_mode();
                 match kani_contract_mode {
                     kani::internal::RECURSION_CHECK => {
                         #recursion_closure;