Fix contract of constant fn with effect feature

kani-compiler/src/kani_middle/stubbing/mod.rs

@@ -17,7 +17,7 @@ use rustc_smir::rustc_internal;
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::visit::{Location, MirVisitor};
 use stable_mir::mir::Constant;
-use stable_mir::ty::FnDef;
+use stable_mir::ty::{FnDef, RigidTy, TyKind};
 use stable_mir::{CrateDef, CrateItem};
 
 use self::annotations::update_stub_mapping;
@@ -37,6 +37,24 @@ pub fn harness_stub_map(
     stub_pairs
 }
 
+/// Retrieve the index of the host parameter if old definition has one, but not the new definition.
+///
+/// This is to allow constant functions to be stubbed by non-constant functions when the
+/// `effect` feature is on.
+///
+/// Note that the opposite is not supported today, but users should be able to change their stubs.
+///
+/// Note that this has no effect at runtime.
+pub fn contract_host_param(tcx: TyCtxt, old_def: FnDef, new_def: FnDef) -> Option<usize> {
+    let old_generics = tcx.generics_of(rustc_internal::internal(tcx, old_def.def_id()));
+    let new_generics = tcx.generics_of(rustc_internal::internal(tcx, new_def.def_id()));
+    if old_generics.host_effect_index.is_some() && new_generics.host_effect_index.is_none() {
+        old_generics.host_effect_index
+    } else {
+        None
+    }
+}
+
 /// Checks whether the stub is compatible with the original function/method: do
 /// the arities and types (of the parameters and return values) match up? This
 /// does **NOT** check whether the type variables are constrained to implement
@@ -61,15 +79,26 @@ pub fn check_compatibility(tcx: TyCtxt, old_def: FnDef, new_def: FnDef) -> Resul
     // Check whether the numbers of generic parameters match.
     let old_def_id = rustc_internal::internal(tcx, old_def.def_id());
     let new_def_id = rustc_internal::internal(tcx, new_def.def_id());
-    let old_num_generics = tcx.generics_of(old_def_id).count();
-    let stub_num_generics = tcx.generics_of(new_def_id).count();
-    if old_num_generics != stub_num_generics {
+    let old_ty = rustc_internal::stable(tcx.type_of(old_def_id)).value;
+    let new_ty = rustc_internal::stable(tcx.type_of(new_def_id)).value;
+    let TyKind::RigidTy(RigidTy::FnDef(_, mut old_args)) = old_ty.kind() else {
+        unreachable!("Expected function, but found {old_ty}")
+    };
+    let TyKind::RigidTy(RigidTy::FnDef(_, new_args)) = new_ty.kind() else {
+        unreachable!("Expected function, but found {new_ty}")
+    };
+    if let Some(idx) = contract_host_param(tcx, old_def, new_def) {
+        old_args.0.remove(idx);
+    }
+
+    // TODO: We should check for the parameter type too or replacement will fail.
+    if old_args.0.len() != new_args.0.len() {
         let msg = format!(
             "mismatch in the number of generic parameters: original function/method `{}` takes {} generic parameters(s), stub `{}` takes {}",
             old_def.name(),
-            old_num_generics,
+            old_args.0.len(),
             new_def.name(),
-            stub_num_generics
+            new_args.0.len(),
         );
         return Err(msg);
     }

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

@@ -3,7 +3,7 @@
 //! This module contains code related to the MIR-to-MIR pass that performs the
 //! stubbing of functions and methods.
 use crate::kani_middle::codegen_units::Stubs;
-use crate::kani_middle::stubbing::validate_stub_const;
+use crate::kani_middle::stubbing::{contract_host_param, validate_stub_const};
 use crate::kani_middle::transform::body::{MutMirVisitor, MutableBody};
 use crate::kani_middle::transform::{TransformPass, TransformationType};
 use crate::kani_queries::QueryDb;
@@ -46,8 +46,12 @@ impl TransformPass for FnStubPass {
     fn transform(&self, tcx: TyCtxt, body: Body, instance: Instance) -> (bool, Body) {
         trace!(function=?instance.name(), "transform");
         let ty = instance.ty();
-        if let TyKind::RigidTy(RigidTy::FnDef(fn_def, args)) = ty.kind() {
+        if let TyKind::RigidTy(RigidTy::FnDef(fn_def, mut args)) = ty.kind() {
             if let Some(replace) = self.stubs.get(&fn_def) {
+                if let Some(idx) = contract_host_param(tcx, fn_def, *replace) {
+                    debug!(?idx, "FnStubPass::transform remove_host_param");
+                    args.0.remove(idx);
+                }
                 let new_instance = Instance::resolve(*replace, &args).unwrap();
                 debug!(from=?instance.name(), to=?new_instance.name(), "FnStubPass::transform");
                 if let Some(body) = FnStubValidator::validate(tcx, (fn_def, *replace), new_instance)

tests/expected/function-contract/const_fn_with_effect.expected

@@ -0,0 +1,2 @@
+Checking harness check...
+VERIFICATION:- SUCCESSFUL

tests/expected/function-contract/const_fn_with_effect.rs

@@ -0,0 +1,18 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts -Zmem-predicates
+
+//! Check that Kani contract can be applied to a constant function.
+//! <https://github.com/model-checking/kani/issues/3258>
+
+#![feature(effects)]
+
+#[kani::requires(kani::mem::can_dereference(arg))]
+const unsafe fn dummy<T>(arg: *const T) -> T {
+    std::ptr::read(arg)
+}
+
+#[kani::proof_for_contract(dummy)]
+fn check() {
+    unsafe { dummy(&kani::any::<u8>()) };
+}

tests/script-based-pre/verify_std_cmd/verify_std.expected

@@ -6,7 +6,10 @@ VERIFICATION:- SUCCESSFUL
 Checking harness verify::harness...
 VERIFICATION:- SUCCESSFUL
 
+Checking harness verify::check_dummy_read...
+VERIFICATION:- SUCCESSFUL
+
 Checking harness num::verify::check_non_zero...
 VERIFICATION:- SUCCESSFUL
 
-Complete - 3 successfully verified harnesses, 0 failures, 3 total.
+Complete - 4 successfully verified harnesses, 0 failures, 4 total.

tests/script-based-pre/verify_std_cmd/verify_std.sh

@@ -47,6 +47,19 @@ pub mod verify {
     fn fake_function(x: bool) -> bool {
         x
     }
+
+    #[kani::proof_for_contract(dummy_read)]
+    fn check_dummy_read() {
+        let val: char = kani::any();
+        assert_eq!(unsafe { dummy_read(&val) }, val);
+    }
+
+    /// Ensure we can verify constant functions.
+    #[kani::requires(kani::mem::can_dereference(ptr))]
+    #[rustc_diagnostic_item = "dummy_read"]
+    const unsafe fn dummy_read<T: Copy>(ptr: *const T) -> T {
+        *ptr
+    }
 }
 '