Introduce test reachability mode (rust-lang#1914)

Author: tedinski

kani-compiler/kani_queries/src/lib.rs

@@ -26,6 +26,8 @@ pub enum ReachabilityType {
     None,
     /// Start the cross-crate reachability analysis from all public functions in the local crate.
     PubFns,
+    /// Start the cross-crate reachability analysis from all *test* (i.e. `#[test]`) harnesses in the local crate.
+    Tests,
 }
 
 impl Default for ReachabilityType {

kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs

@@ -284,6 +284,38 @@ impl<'tcx> GotocCtx<'tcx> {
         }
     }
 
+    /// Does this `def_id` have `#[rustc_test_marker]`?
+    pub fn is_test_harness_description(&self, def_id: DefId) -> bool {
+        let attrs = self.tcx.get_attrs_unchecked(def_id);
+
+        self.tcx.sess.contains_name(attrs, rustc_span::symbol::sym::rustc_test_marker)
+    }
+    /// Is this the closure inside of a test description const (i.e. macro expanded from a `#[test]`)?
+    ///
+    /// We're trying to detect the closure (`||`) inside code like:
+    ///
+    /// ```ignore
+    /// #[rustc_test_marker]
+    /// pub const check_2: test::TestDescAndFn = test::TestDescAndFn {
+    ///     desc: ...,
+    ///     testfn: test::StaticTestFn(|| test::assert_test_result(check_2())),
+    /// };
+    /// ```
+    pub fn is_test_harness_closure(&self, def_id: DefId) -> bool {
+        if !def_id.is_local() {
+            return false;
+        }
+
+        let local_def_id = def_id.expect_local();
+        let hir_id = self.tcx.hir().local_def_id_to_hir_id(local_def_id);
+
+        // The parent item of the closure appears to reliably be the `const` declaration item.
+        let parent_id = self.tcx.hir().get_parent_item(hir_id);
+        let parent_def_id = parent_id.to_def_id();
+
+        self.is_test_harness_description(parent_def_id)
+    }
+
     /// We record test harness information in kani-metadata, just like we record
     /// proof harness information. This is used to support e.g. cargo-kani assess.
     ///
@@ -293,34 +325,16 @@ impl<'tcx> GotocCtx<'tcx> {
     /// as it add asserts for tests that return `Result` types.
     fn record_test_harness_metadata(&mut self) {
         let def_id = self.current_fn().instance().def_id();
-        if def_id.is_local() {
-            let local_def_id = def_id.expect_local();
-            let hir_id = self.tcx.hir().local_def_id_to_hir_id(local_def_id);
-
-            // We want to detect the case where we're codegen'ing the closure inside what test "descriptions"
-            // are macro-expanded to:
-            //
-            // #[rustc_test_marker]
-            // pub const check_2: test::TestDescAndFn = test::TestDescAndFn {
-            //     desc: ...,
-            //     testfn: test::StaticTestFn(|| test::assert_test_result(check_2())),
-            // };
-
-            // The parent item of the closure appears to reliably be the `const` declaration item.
-            let parent_id = self.tcx.hir().get_parent_item(hir_id);
-            let attrs = self.tcx.get_attrs_unchecked(parent_id.to_def_id());
-
-            if self.tcx.sess.contains_name(attrs, rustc_span::symbol::sym::rustc_test_marker) {
-                let loc = self.codegen_span(&self.current_fn().mir().span);
-                self.test_harnesses.push(HarnessMetadata {
-                    pretty_name: self.current_fn().readable_name().to_owned(),
-                    mangled_name: self.current_fn().name(),
-                    original_file: loc.filename().unwrap(),
-                    original_start_line: loc.start_line().unwrap() as usize,
-                    original_end_line: loc.end_line().unwrap() as usize,
-                    unwind_value: None,
-                })
-            }
+        if self.is_test_harness_closure(def_id) {
+            let loc = self.codegen_span(&self.current_fn().mir().span);
+            self.test_harnesses.push(HarnessMetadata {
+                pretty_name: self.current_fn().readable_name().to_owned(),
+                mangled_name: self.current_fn().name(),
+                original_file: loc.filename().unwrap(),
+                original_start_line: loc.start_line().unwrap() as usize,
+                original_end_line: loc.end_line().unwrap() as usize,
+                unwind_value: None,
+            })
         }
     }
 

kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs

@@ -5,7 +5,9 @@
 
 use crate::codegen_cprover_gotoc::GotocCtx;
 use crate::kani_middle::mir_transform;
-use crate::kani_middle::reachability::{collect_reachable_items, filter_crate_items};
+use crate::kani_middle::reachability::{
+    collect_reachable_items, filter_closures_in_const_crate_items, filter_crate_items,
+};
 use bitflags::_core::any::Any;
 use cbmc::goto_program::{symtab_transformer, Location};
 use cbmc::{InternedString, MachineModel};
@@ -392,6 +394,14 @@ fn collect_codegen_items<'tcx>(gcx: &GotocCtx<'tcx>) -> Vec<MonoItem<'tcx>> {
             let harnesses = filter_crate_items(tcx, |_, def_id| gcx.is_proof_harness(def_id));
             collect_reachable_items(tcx, &harnesses).into_iter().collect()
         }
+        ReachabilityType::Tests => {
+            // We're iterating over crate items here, so what we have to codegen is the "test description" containing the
+            // test closure that we want to execute
+            let harnesses = filter_closures_in_const_crate_items(tcx, |_, def_id| {
+                gcx.is_test_harness_description(def_id)
+            });
+            collect_reachable_items(tcx, &harnesses).into_iter().collect()
+        }
         ReachabilityType::None => Vec::new(),
         ReachabilityType::PubFns => {
             let entry_fn = tcx.entry_fn(()).map(|(id, _)| id);

kani-compiler/src/kani_middle/reachability.rs

@@ -18,6 +18,7 @@ use tracing::{debug, debug_span, trace, warn};
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
+use rustc_hir::def_id::DefId;
 use rustc_middle::mir::interpret::{AllocId, ConstValue, ErrorHandled, GlobalAlloc, Scalar};
 use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::mir::visit::Visitor as MirVisitor;
@@ -28,9 +29,8 @@ use rustc_middle::span_bug;
 use rustc_middle::ty::adjustment::PointerCast;
 use rustc_middle::ty::{
     Closure, ClosureKind, ConstKind, Instance, InstanceDef, ParamEnv, Ty, TyCtxt, TyKind,
-    TypeFoldable, VtblEntry,
+    TypeFoldable, VtblEntry, WithOptConstParam,
 };
-use rustc_span::def_id::DefId;
 
 use crate::kani_middle::coercion;
 
@@ -54,12 +54,11 @@ pub fn collect_reachable_items<'tcx>(
     sorted_items
 }
 
-/// Collect all items in the crate that matches the given predicate.
+/// Collect all (top-level) items in the crate that matches the given predicate.
 pub fn filter_crate_items<F>(tcx: TyCtxt, predicate: F) -> Vec<MonoItem>
 where
     F: FnMut(TyCtxt, DefId) -> bool,
 {
-    // Filter proof harnesses.
     let mut filter = predicate;
     tcx.hir_crate_items(())
         .items()
@@ -70,6 +69,31 @@ where
         .collect()
 }
 
+/// Use a predicate to find `const` declarations, then extract all closures from those declarations
+///
+/// Probably only specifically useful with a predicate to find `TestDescAndFn` const declarations from
+/// tests and extract the closures from them.
+pub fn filter_closures_in_const_crate_items<F>(tcx: TyCtxt, mut predicate: F) -> Vec<MonoItem>
+where
+    F: FnMut(TyCtxt, DefId) -> bool,
+{
+    let mut roots = Vec::new();
+    for hir_id in tcx.hir_crate_items(()).items() {
+        let def_id = hir_id.owner_id.def_id.to_def_id();
+        if predicate(tcx, def_id) {
+            // The predicate should only ever apply to monomorphic items
+            let instance = Instance::mono(tcx, def_id);
+            let body = tcx.instance_mir(InstanceDef::Item(WithOptConstParam::unknown(def_id)));
+            let mut extrator =
+                ConstMonoItemExtractor { tcx, body, instance, collected: FxHashSet::default() };
+            extrator.visit_body(body);
+
+            roots.extend(extrator.collected);
+        }
+    }
+    roots
+}
+
 struct MonoItemsCollector<'tcx> {
     /// The compiler context.
     tcx: TyCtxt<'tcx>,
@@ -477,3 +501,50 @@ fn collect_alloc_items<'tcx>(tcx: TyCtxt<'tcx>, alloc_id: AllocId) -> Vec<MonoIt
     };
     items
 }
+
+/// This MIR Visitor is intended for one specific purpose:
+/// Find the closure that exist inside a top-level const declaration generated by
+/// test declarations. This allows us to treat this closure instance as a root for
+/// the reachability analysis.
+///
+/// Entry into this visitor will be via `visit_body`
+struct ConstMonoItemExtractor<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    collected: FxHashSet<MonoItem<'tcx>>,
+    instance: Instance<'tcx>,
+    body: &'a Body<'tcx>,
+}
+
+impl<'a, 'tcx> MirVisitor<'tcx> for ConstMonoItemExtractor<'a, 'tcx> {
+    #[allow(clippy::single_match)]
+    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
+        trace!(rvalue=?*rvalue, "visit_rvalue");
+
+        match *rvalue {
+            Rvalue::Cast(CastKind::Pointer(PointerCast::ClosureFnPointer(_)), ref operand, _) => {
+                let source_ty = operand.ty(self.body, self.tcx);
+                let source_ty = self.instance.subst_mir_and_normalize_erasing_regions(
+                    self.tcx,
+                    ParamEnv::reveal_all(),
+                    source_ty,
+                );
+                match *source_ty.kind() {
+                    Closure(def_id, substs) => {
+                        let instance = Instance::resolve_closure(
+                            self.tcx,
+                            def_id,
+                            substs,
+                            ClosureKind::FnOnce,
+                        )
+                        .expect("failed to normalize and resolve closure during codegen");
+                        self.collected.insert(MonoItem::Fn(instance.polymorphize(self.tcx)));
+                    }
+                    _ => unreachable!("Unexpected type: {:?}", source_ty),
+                }
+            }
+            _ => { /* not interesting */ }
+        }
+
+        self.super_rvalue(rvalue, location);
+    }
+}

kani-driver/src/assess.rs

@@ -216,7 +216,7 @@ pub(crate) fn cargokani_assess_main(mut ctx: KaniSession) -> Result<()> {
     ctx.args.unwind = Some(1);
     ctx.args.tests = true;
     ctx.args.output_format = crate::args::OutputFormat::Terse;
-    ctx.codegen_pub_fns = true;
+    ctx.codegen_tests = true;
     if ctx.args.jobs.is_none() {
         // assess will default to fully parallel instead of single-threaded.
         // can be overridden with e.g. `cargo kani --enable-unstable -j 8 assess`

kani-driver/src/call_cargo.rs

@@ -82,6 +82,9 @@ impl KaniSession {
             ReachabilityMode::AllPubFns => {
                 pkg_args.extend(["--".into(), "--reachability=pub_fns".into()]);
             }
+            ReachabilityMode::Tests => {
+                pkg_args.extend(["--".into(), "--reachability=tests".into()]);
+            }
         }
 
         // Only joing them at the end. All kani flags must come first.

kani-driver/src/call_single_file.rs

@@ -53,6 +53,7 @@ impl KaniSession {
                 ReachabilityMode::Legacy => "--reachability=legacy",
                 ReachabilityMode::ProofHarnesses => "--reachability=harnesses",
                 ReachabilityMode::AllPubFns => "--reachability=pub_fns",
+                ReachabilityMode::Tests => "--reachability=tests",
             }
             .into(),
         );

kani-driver/src/session.rs

@@ -26,7 +26,7 @@ pub struct KaniSession {
     /// Include all publicly-visible symbols in the generated goto binary, not just those reachable from
     /// a proof harness. Useful when attempting to verify things that were not annotated with kani
     /// proof attributes.
-    pub codegen_pub_fns: bool,
+    pub codegen_tests: bool,
 
     /// The location we found the 'kani_rustc' command
     pub kani_compiler: PathBuf,
@@ -56,7 +56,7 @@ impl KaniSession {
 
         Ok(KaniSession {
             args,
-            codegen_pub_fns: false,
+            codegen_tests: false,
             kani_compiler: install.kani_compiler()?,
             kani_lib_c: install.kani_lib_c()?,
             kani_c_stubs: install.kani_c_stubs()?,
@@ -76,7 +76,9 @@ impl KaniSession {
     pub fn reachability_mode(&self) -> ReachabilityMode {
         if self.args.legacy_linker {
             ReachabilityMode::Legacy
-        } else if self.args.function.is_some() || self.codegen_pub_fns {
+        } else if self.codegen_tests {
+            ReachabilityMode::Tests
+        } else if self.args.function.is_some() {
             ReachabilityMode::AllPubFns
         } else {
             ReachabilityMode::ProofHarnesses
@@ -88,6 +90,7 @@ pub enum ReachabilityMode {
     Legacy,
     ProofHarnesses,
     AllPubFns,
+    Tests,
 }
 
 impl Drop for KaniSession {