Automatic Derivation Fixes (#4194)

Fixes and improvements to the derivation of Arbitrary in the compiler
introduced in #4167, along with a fix for #4189 and some other small
improvements to autoharness that I noticed along the way. Best reviewed
commit by commit.

Resolves #4189
Towards #3832 

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

Author: carolynzech

kani-compiler/src/kani_middle/codegen_units.rs

@@ -18,7 +18,7 @@ use crate::kani_middle::resolve::expect_resolve_fn;
 use crate::kani_middle::stubbing::{check_compatibility, harness_stub_map};
 use crate::kani_middle::{can_derive_arbitrary, implements_arbitrary};
 use crate::kani_queries::QueryDb;
-use fxhash::FxHashMap;
+use fxhash::{FxHashMap, FxHashSet};
 use kani_metadata::{
     ArtifactType, AssignsContract, AutoHarnessMetadata, AutoHarnessSkipReason, HarnessKind,
     HarnessMetadata, KaniMetadata,
@@ -390,8 +390,16 @@ fn automatic_harness_partition(
     crate_name: &str,
     kani_any_def: FnDef,
 ) -> (Vec<Instance>, BTreeMap<String, AutoHarnessSkipReason>) {
-    let crate_fns =
-        stable_mir::all_local_items().into_iter().filter(|item| item.ty().kind().is_fn());
+    let crate_fn_defs = stable_mir::local_crate().fn_defs().into_iter().collect::<FxHashSet<_>>();
+    // Filter out CrateItems that are functions, but not functions defined in the crate itself, i.e., rustc-inserted functions
+    // (c.f. https://github.com/model-checking/kani/issues/4189)
+    let crate_fns = stable_mir::all_local_items().into_iter().filter(|item| {
+        if let TyKind::RigidTy(RigidTy::FnDef(def, _)) = item.ty().kind() {
+            crate_fn_defs.contains(&def)
+        } else {
+            false
+        }
+    });
 
     let included_set = make_regex_set(args.autoharness_included_patterns.clone());
     let excluded_set = make_regex_set(args.autoharness_excluded_patterns.clone());
@@ -437,12 +445,15 @@ fn automatic_harness_partition(
         // Note that we've already filtered out generic functions, so we know that each of these arguments has a concrete type.
         let mut problematic_args = vec![];
         for (idx, arg) in body.arg_locals().iter().enumerate() {
-            let implements_arbitrary = ty_arbitrary_cache.entry(arg.ty).or_insert_with(|| {
-                implements_arbitrary(arg.ty, kani_any_def)
-                    || can_derive_arbitrary(arg.ty, kani_any_def)
-            });
+            if !ty_arbitrary_cache.contains_key(&arg.ty) {
+                let impls_arbitrary =
+                    implements_arbitrary(arg.ty, kani_any_def, &mut ty_arbitrary_cache)
+                        || can_derive_arbitrary(arg.ty, kani_any_def, &mut ty_arbitrary_cache);
+                ty_arbitrary_cache.insert(arg.ty, impls_arbitrary);
+            }
+            let impls_arbitrary = ty_arbitrary_cache.get(&arg.ty).unwrap();
 
-            if !(*implements_arbitrary) {
+            if !impls_arbitrary {
                 // Find the name of the argument by referencing var_debug_info.
                 // Note that enumerate() starts at 0, while StableMIR argument_index starts at 1, hence the idx+1.
                 let arg_name = body

kani-compiler/src/kani_middle/mod.rs

@@ -6,6 +6,7 @@
 use std::collections::HashSet;
 
 use crate::kani_queries::QueryDb;
+use fxhash::FxHashMap;
 use rustc_hir::{def::DefKind, def_id::DefId as InternalDefId, def_id::LOCAL_CRATE};
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
@@ -178,7 +179,19 @@ pub fn stable_fn_def(tcx: TyCtxt, def_id: InternalDefId) -> Option<FnDef> {
 /// ```
 /// So we select the terminator that calls T::kani::Arbitrary::any(), then try to resolve it to an Instance.
 /// `T` implements Arbitrary iff we successfully resolve the Instance.
-fn implements_arbitrary(ty: Ty, kani_any_def: FnDef) -> bool {
+fn implements_arbitrary(
+    ty: Ty,
+    kani_any_def: FnDef,
+    ty_arbitrary_cache: &mut FxHashMap<Ty, bool>,
+) -> bool {
+    if let Some(v) = ty_arbitrary_cache.get(&ty) {
+        return *v;
+    }
+
+    if ty.kind().rigid().is_none() {
+        return false;
+    }
+
     let kani_any_body =
         Instance::resolve(kani_any_def, &GenericArgs(vec![GenericArgKind::Type(ty)]))
             .unwrap()
@@ -192,20 +205,32 @@ fn implements_arbitrary(ty: Ty, kani_any_def: FnDef) -> bool {
         if let TyKind::RigidTy(RigidTy::FnDef(def, args)) =
             func.ty(kani_any_body.arg_locals()).unwrap().kind()
         {
-            return Instance::resolve(def, &args).is_ok();
+            let res = Instance::resolve(def, &args).is_ok();
+            ty_arbitrary_cache.insert(ty, res);
+            return res;
         }
     }
     false
 }
 
 /// Is `ty` a struct or enum whose fields/variants implement Arbitrary?
-fn can_derive_arbitrary(ty: Ty, kani_any_def: FnDef) -> bool {
-    let variants_can_derive = |def: AdtDef| {
+fn can_derive_arbitrary(
+    ty: Ty,
+    kani_any_def: FnDef,
+    ty_arbitrary_cache: &mut FxHashMap<Ty, bool>,
+) -> bool {
+    let mut variants_can_derive = |def: AdtDef, args: GenericArgs| {
         for variant in def.variants_iter() {
             let fields = variant.fields();
             let mut fields_impl_arbitrary = true;
-            for ty in fields.iter().map(|field| field.ty()) {
-                fields_impl_arbitrary &= implements_arbitrary(ty, kani_any_def);
+            for ty in fields.iter().map(|field| field.ty_with_args(&args)) {
+                if let TyKind::RigidTy(RigidTy::Adt(..)) = ty.kind() {
+                    fields_impl_arbitrary &=
+                        can_derive_arbitrary(ty, kani_any_def, ty_arbitrary_cache);
+                } else {
+                    fields_impl_arbitrary &=
+                        implements_arbitrary(ty, kani_any_def, ty_arbitrary_cache);
+                }
             }
             if !fields_impl_arbitrary {
                 return false;
@@ -214,16 +239,22 @@ fn can_derive_arbitrary(ty: Ty, kani_any_def: FnDef) -> bool {
         true
     };
 
-    if let TyKind::RigidTy(RigidTy::Adt(def, _)) = ty.kind() {
+    if let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() {
+        for arg in &args.0 {
+            if let GenericArgKind::Lifetime(..) = arg {
+                return false;
+            }
+        }
+
         match def.kind() {
             AdtKind::Enum => {
                 // Enums with no variants cannot be instantiated
                 if def.num_variants() == 0 {
                     return false;
                 }
-                variants_can_derive(def)
+                variants_can_derive(def, args)
             }
-            AdtKind::Struct => variants_can_derive(def),
+            AdtKind::Struct => variants_can_derive(def, args),
             AdtKind::Union => false,
         }
     } else {

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

@@ -14,6 +14,7 @@ use crate::kani_middle::kani_functions::{KaniHook, KaniIntrinsic, KaniModel};
 use crate::kani_middle::transform::body::{InsertPosition, MutableBody, SourceInstruction};
 use crate::kani_middle::transform::{TransformPass, TransformationType};
 use crate::kani_queries::QueryDb;
+use fxhash::FxHashMap;
 use rustc_middle::ty::TyCtxt;
 use stable_mir::CrateDef;
 use stable_mir::mir::mono::Instance;
@@ -109,14 +110,14 @@ impl TransformPass for AutomaticArbitraryPass {
         let binding = instance.args();
         let ty = binding.0[0].expect_ty();
 
-        if implements_arbitrary(*ty, self.kani_any) {
+        if implements_arbitrary(*ty, self.kani_any, &mut FxHashMap::default()) {
             return (false, body);
         }
 
-        if let TyKind::RigidTy(RigidTy::Adt(def, ..)) = ty.kind() {
+        if let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() {
             match def.kind() {
-                AdtKind::Enum => (true, self.generate_enum_body(def, body)),
-                AdtKind::Struct => (true, self.generate_struct_body(def, body)),
+                AdtKind::Enum => (true, self.generate_enum_body(def, args, body)),
+                AdtKind::Struct => (true, self.generate_struct_body(def, args, body)),
                 AdtKind::Union => unexpected_ty(ty),
             }
         } else {
@@ -151,7 +152,8 @@ impl AutomaticArbitraryPass {
     /// This function will panic if a field type does not implement Arbitrary.
     fn call_kani_any_for_variant(
         &self,
-        def: AdtDef,
+        adt_def: AdtDef,
+        adt_args: &GenericArgs,
         body: &mut MutableBody,
         source: &mut SourceInstruction,
         variant: VariantDef,
@@ -160,7 +162,7 @@ impl AutomaticArbitraryPass {
         let mut field_locals = vec![];
 
         // Construct nondeterministic values for each of the variant's fields
-        for ty in fields.iter().map(|field| field.ty()) {
+        for ty in fields.iter().map(|field| field.ty_with_args(adt_args)) {
             let lcl = self.call_kani_any_for_ty(body, ty, source);
             field_locals.push(lcl);
         }
@@ -173,7 +175,7 @@ impl AutomaticArbitraryPass {
         );
         let mut assign_instr = SourceInstruction::Terminator { bb: source.bb() - 1 };
         let rvalue = Rvalue::Aggregate(
-            AggregateKind::Adt(def, variant.idx, GenericArgs(vec![]), None, None),
+            AggregateKind::Adt(adt_def, variant.idx, adt_args.clone(), None, None),
             field_locals.into_iter().map(|lcl| Operand::Move(lcl.into())).collect(),
         );
         body.assign_to(Place::from(0), rvalue, &mut assign_instr, InsertPosition::Before);
@@ -193,7 +195,7 @@ impl AutomaticArbitraryPass {
     ///   _ => Enum::LastVariant
     /// }
     /// ```
-    fn generate_enum_body(&self, def: AdtDef, body: Body) -> Body {
+    fn generate_enum_body(&self, def: AdtDef, args: GenericArgs, body: Body) -> Body {
         // Autoharness only deems a function with an enum eligible if it has at least one variant, c.f. `can_derive_arbitrary`
         assert!(def.num_variants() > 0);
 
@@ -220,7 +222,7 @@ impl AutomaticArbitraryPass {
         let mut branches: Vec<(u128, BasicBlockIdx)> = vec![];
         for variant in def.variants_iter() {
             let target_bb =
-                self.call_kani_any_for_variant(def, &mut new_body, &mut source, variant);
+                self.call_kani_any_for_variant(def, &args, &mut new_body, &mut source, variant);
             branches.push((variant.idx.to_index() as u128, target_bb));
         }
 
@@ -246,53 +248,37 @@ impl AutomaticArbitraryPass {
     ///   ...
     /// }
     /// ```
-    fn generate_struct_body(&self, def: AdtDef, body: Body) -> Body {
+    fn generate_struct_body(&self, def: AdtDef, args: GenericArgs, body: Body) -> Body {
         assert_eq!(def.num_variants(), 1);
 
         let mut new_body = MutableBody::from(body);
         new_body.clear_body(TerminatorKind::Unreachable);
         let mut source = SourceInstruction::Terminator { bb: 0 };
 
         let variant = def.variants()[0];
-        self.call_kani_any_for_variant(def, &mut new_body, &mut source, variant);
+        self.call_kani_any_for_variant(def, &args, &mut new_body, &mut source, variant);
 
         new_body.into()
     }
 }
 /// Transform the dummy body of an automatic_harness Kani intrinsic to be a proof harness for a given function.
 #[derive(Debug)]
 pub struct AutomaticHarnessPass {
-    /// The FnDef of KaniModel::Any
     kani_any: FnDef,
     init_contracts_hook: Instance,
-    /// All of the automatic harness Instances that we generated in the CodegenUnits constructor
-    automatic_harnesses: Vec<Instance>,
+    kani_autoharness_intrinsic: FnDef,
 }
 
 impl AutomaticHarnessPass {
-    // FIXME: this is a bit clunky.
-    // Historically, in codegen_crate, we reset the BodyTransformation cache on a per-unit basis,
-    // so the BodyTransformation constructor only accepts a CodegenUnit and thus this constructor can only accept a unit.
-    // Later, we changed codegen to reset the cache on a per-harness basis (for uninitialized memory instrumentation).
-    // So BodyTransformation should really be changed to reflect that, so that this constructor can just save the one automatic harness it should transform
-    // and not all of the possibilities.
-    pub fn new(unit: &CodegenUnit, query_db: &QueryDb) -> Self {
+    pub fn new(query_db: &QueryDb) -> Self {
         let kani_fns = query_db.kani_functions();
-        let harness_intrinsic = *kani_fns.get(&KaniIntrinsic::AutomaticHarness.into()).unwrap();
+        let kani_autoharness_intrinsic =
+            *kani_fns.get(&KaniIntrinsic::AutomaticHarness.into()).unwrap();
         let kani_any = *kani_fns.get(&KaniModel::Any.into()).unwrap();
         let init_contracts_hook = *kani_fns.get(&KaniHook::InitContracts.into()).unwrap();
         let init_contracts_hook =
             Instance::resolve(init_contracts_hook, &GenericArgs(vec![])).unwrap();
-        let automatic_harnesses = unit
-            .harnesses
-            .iter()
-            .cloned()
-            .filter(|harness| {
-                let (def, _) = harness.ty().kind().fn_def().unwrap();
-                def == harness_intrinsic
-            })
-            .collect::<Vec<_>>();
-        Self { kani_any, init_contracts_hook, automatic_harnesses }
+        Self { kani_any, init_contracts_hook, kani_autoharness_intrinsic }
     }
 }
 
@@ -314,7 +300,7 @@ impl TransformPass for AutomaticHarnessPass {
     fn transform(&mut self, tcx: TyCtxt, body: Body, instance: Instance) -> (bool, Body) {
         debug!(function=?instance.name(), "AutomaticHarnessPass::transform");
 
-        if !self.automatic_harnesses.contains(&instance) {
+        if instance.def.def_id() != self.kani_autoharness_intrinsic.def_id() {
             return (false, body);
         }
 

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

@@ -75,7 +75,7 @@ impl BodyTransformation {
         let safety_check_type = CheckType::new_safety_check_assert_assume(queries);
         let unsupported_check_type = CheckType::new_unsupported_check_assert_assume_false(queries);
         // This has to come first, since creating harnesses affects later stubbing and contract passes.
-        transformer.add_pass(queries, AutomaticHarnessPass::new(unit, queries));
+        transformer.add_pass(queries, AutomaticHarnessPass::new(queries));
         transformer.add_pass(queries, AutomaticArbitraryPass::new(unit, queries));
         transformer.add_pass(queries, FnStubPass::new(&unit.stubs));
         transformer.add_pass(queries, ExternFnStubPass::new(&unit.stubs));