Autoharness: Derive Arbitrary for structs and enums

Cargo.lock

@@ -1084,6 +1084,7 @@ dependencies = [
  "charon",
  "clap",
  "cprover_bindings",
+ "fxhash",
  "itertools 0.14.0",
  "kani_metadata",
  "lazy_static",

docs/src/reference/experimental/autoharness.md

@@ -103,10 +103,10 @@ please add them to [this GitHub issue](https://github.com/model-checking/kani/is
 
 ## Limitations
 ### Arguments Implementing Arbitrary
-Kani will only generate an automatic harness for a function if it can determine that all of the function's arguments implement Arbitrary.
-It does not attempt to derive/implement Arbitrary for any types, even if those types could implement Arbitrary.
-For example, imagine a user defines `struct MyStruct { x: u8, y: u8}`, but does not derive or implement Arbitrary for `MyStruct`.
-Kani does not attempt to add such derivations itself, so it will not generate a harness for a function that takes `MyStruct` as input.
+Kani will only generate an automatic harness for a function if it can represent each of its arguments nondeterministically, without bounds.
+In technical terms, each of the arguments needs to implement the `Arbitrary` trait or be capable of deriving it.
+Kani will detect if a struct or enum could implement `Arbitrary` and derive it automatically.
+Note that this automatic derivation feature is only available for autoharness.
 
 ### Generic Functions
 The current implementation does not generate harnesses for generic functions.

kani-compiler/Cargo.toml

@@ -12,6 +12,7 @@ publish = false
 cbmc = { path = "../cprover_bindings", package = "cprover_bindings", optional = true }
 charon = { path = "../charon/charon", optional = true, default-features = false } 
 clap = { version = "4.4.11", features = ["derive", "cargo"] }
+fxhash = "0.2.1"
 itertools = "0.14"
 kani_metadata = { path = "../kani_metadata" }
 lazy_static = "1.5.0"

kani-compiler/src/kani_middle/codegen_units.rs

@@ -16,7 +16,9 @@ use crate::kani_middle::metadata::{
 use crate::kani_middle::reachability::filter_crate_items;
 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 kani_metadata::{
     ArtifactType, AssignsContract, AutoHarnessMetadata, AutoHarnessSkipReason, HarnessKind,
     HarnessMetadata, KaniMetadata,
@@ -26,8 +28,8 @@ use rustc_hir::def_id::DefId;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::OutputType;
 use rustc_smir::rustc_internal;
-use stable_mir::mir::{TerminatorKind, mono::Instance};
-use stable_mir::ty::{FnDef, GenericArgKind, GenericArgs, IndexedVal, RigidTy, TyKind};
+use stable_mir::mir::mono::Instance;
+use stable_mir::ty::{FnDef, GenericArgKind, GenericArgs, IndexedVal, RigidTy, Ty, TyKind};
 use stable_mir::{CrateDef, CrateItem};
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::fs::File;
@@ -394,10 +396,13 @@ fn automatic_harness_partition(
     let included_set = make_regex_set(args.autoharness_included_patterns.clone());
     let excluded_set = make_regex_set(args.autoharness_excluded_patterns.clone());
 
+    // Cache whether a type implements or can derive Arbitrary
+    let mut ty_arbitrary_cache: FxHashMap<Ty, bool> = FxHashMap::default();
+
     // If `func` is not eligible for an automatic harness, return the reason why; if it is eligible, return None.
     // Note that we only return one reason for ineligiblity, when there could be multiple;
     // we can revisit this implementation choice in the future if users request more verbose output.
-    let skip_reason = |fn_item: CrateItem| -> Option<AutoHarnessSkipReason> {
+    let mut skip_reason = |fn_item: CrateItem| -> Option<AutoHarnessSkipReason> {
         if KaniAttributes::for_def_id(tcx, fn_item.def_id()).is_kani_instrumentation() {
             return Some(AutoHarnessSkipReason::KaniImpl);
         }
@@ -432,25 +437,12 @@ 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 kani_any_body =
-                Instance::resolve(kani_any_def, &GenericArgs(vec![GenericArgKind::Type(arg.ty)]))
-                    .unwrap()
-                    .body()
-                    .unwrap();
-
-            let implements_arbitrary = if let TerminatorKind::Call { func, .. } =
-                &kani_any_body.blocks[0].terminator.kind
-            {
-                if let Some((def, args)) = func.ty(body.arg_locals()).unwrap().kind().fn_def() {
-                    Instance::resolve(def, args).is_ok()
-                } else {
-                    false
-                }
-            } else {
-                false
-            };
-
-            if !implements_arbitrary {
+            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 !(*implements_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

@@ -9,8 +9,11 @@ use crate::kani_queries::QueryDb;
 use rustc_hir::{def::DefKind, def_id::DefId as InternalDefId, def_id::LOCAL_CRATE};
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
-use stable_mir::mir::mono::MonoItem;
-use stable_mir::ty::{FnDef, RigidTy, Span as SpanStable, Ty, TyKind};
+use stable_mir::mir::TerminatorKind;
+use stable_mir::mir::mono::{Instance, MonoItem};
+use stable_mir::ty::{
+    AdtDef, AdtKind, FnDef, GenericArgKind, GenericArgs, RigidTy, Span as SpanStable, Ty, TyKind,
+};
 use stable_mir::visitor::{Visitable, Visitor as TyVisitor};
 use stable_mir::{CrateDef, DefId};
 use std::ops::ControlFlow;
@@ -165,3 +168,65 @@ pub fn stable_fn_def(tcx: TyCtxt, def_id: InternalDefId) -> Option<FnDef> {
         None
     }
 }
+
+/// Inspect a `kani::any<T>()` call to determine if `T: Arbitrary`
+/// `kani_any_def` refers to a function that looks like:
+/// ```rust
+/// fn any<T: Arbitrary>() -> T {
+///   T::any()
+/// }
+/// ```
+/// 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 {
+    let kani_any_body =
+        Instance::resolve(kani_any_def, &GenericArgs(vec![GenericArgKind::Type(ty)]))
+            .unwrap()
+            .body()
+            .unwrap();
+
+    for bb in kani_any_body.blocks.iter() {
+        let TerminatorKind::Call { func, .. } = &bb.terminator.kind else {
+            continue;
+        };
+        if let TyKind::RigidTy(RigidTy::FnDef(def, args)) =
+            func.ty(kani_any_body.arg_locals()).unwrap().kind()
+        {
+            return Instance::resolve(def, &args).is_ok();
+        }
+    }
+    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| {
+        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);
+            }
+            if !fields_impl_arbitrary {
+                return false;
+            }
+        }
+        true
+    };
+
+    if let TyKind::RigidTy(RigidTy::Adt(def, _)) = ty.kind() {
+        match def.kind() {
+            AdtKind::Enum => {
+                // Enums with no variants cannot be instantiated
+                if def.num_variants() == 0 {
+                    return false;
+                }
+                variants_can_derive(def)
+            }
+            AdtKind::Struct => variants_can_derive(def),
+            AdtKind::Union => false,
+        }
+    } else {
+        false
+    }
+}