Compile harnesses with the same set of stubs together (rust-lang#2514)

Improve Kani compiler to compile harnesses that have exact the same set of stubs in a single Rust compiler session. This reduces the number of times we compile a crate when it has multiple harnesses use the same set of stubs.

Co-authored-by: Zyad Hassan <88045115+zhassan-aws@users.noreply.github.com>

Author: celinval

kani-compiler/src/kani_compiler.rs

@@ -33,7 +33,7 @@ use rustc_interface::Config;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::{ErrorOutputType, OutputType};
 use rustc_span::ErrorGuaranteed;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::fs::File;
 use std::io::BufWriter;
 use std::mem;
@@ -67,7 +67,7 @@ fn backend(queries: Arc<Mutex<QueryDb>>) -> Box<CodegenBackend> {
 type HarnessId = DefPathHash;
 
 /// A set of stubs.
-type Stubs = HashMap<DefPathHash, DefPathHash>;
+type Stubs = BTreeMap<DefPathHash, DefPathHash>;
 
 #[derive(Debug)]
 struct HarnessInfo {
@@ -109,14 +109,15 @@ enum CompilationStage {
     /// Stage where the compiler will perform codegen of all harnesses that don't use stub.
     CodegenNoStubs {
         target_harnesses: Vec<HarnessId>,
-        next_harnesses: Vec<HarnessId>,
+        next_harnesses: Vec<Vec<HarnessId>>,
         all_harnesses: HashMap<HarnessId, HarnessInfo>,
     },
-    /// Stage where the compiler will codegen harnesses that use stub, one at a time.
-    /// Note: We could potentially codegen all harnesses that have the same list of stubs.
+    /// Stage where the compiler will codegen harnesses that use stub, one group at a time.
+    /// The harnesses at this stage are grouped according to the stubs they are using. For now,
+    /// we ensure they have the exact same set of stubs.
     CodegenWithStubs {
-        target_harness: HarnessId,
-        next_harnesses: Vec<HarnessId>,
+        target_harnesses: Vec<HarnessId>,
+        next_harnesses: Vec<Vec<HarnessId>>,
         all_harnesses: HashMap<HarnessId, HarnessInfo>,
     },
     Done,
@@ -166,7 +167,9 @@ impl KaniCompiler {
                 CompilationStage::CodegenNoStubs { .. } => {
                     unreachable!("This stage should always run in the same session as Init");
                 }
-                CompilationStage::CodegenWithStubs { target_harness, all_harnesses, .. } => {
+                CompilationStage::CodegenWithStubs { target_harnesses, all_harnesses, .. } => {
+                    assert!(!target_harnesses.is_empty(), "expected at least one target harness");
+                    let target_harness = &target_harnesses[0];
                     let extra_arg =
                         stubbing::mk_rustc_arg(&all_harnesses[&target_harness].stub_map);
                     let mut args = orig_args.clone();
@@ -192,9 +195,10 @@ impl KaniCompiler {
             }
             CompilationStage::CodegenNoStubs { next_harnesses, all_harnesses, .. }
             | CompilationStage::CodegenWithStubs { next_harnesses, all_harnesses, .. } => {
-                if let Some(target_harness) = next_harnesses.pop() {
+                if let Some(target_harnesses) = next_harnesses.pop() {
+                    assert!(!target_harnesses.is_empty(), "expected at least one target harness");
                     CompilationStage::CodegenWithStubs {
-                        target_harness,
+                        target_harnesses,
                         next_harnesses: mem::take(next_harnesses),
                         all_harnesses: mem::take(all_harnesses),
                     }
@@ -210,6 +214,7 @@ impl KaniCompiler {
 
     /// Run the Rust compiler with the given arguments and pass `&mut self` to handle callbacks.
     fn run_compilation_session(&mut self, args: &[String]) -> Result<(), ErrorGuaranteed> {
+        debug!(?args, "run_compilation_session");
         let queries = self.queries.clone();
         let mut compiler = RunCompiler::new(args, self);
         compiler.set_make_codegen_backend(Some(Box::new(move |_cfg| backend(queries))));
@@ -249,7 +254,7 @@ impl KaniCompiler {
             // Even if no_stubs is empty we still need to store metadata.
             CompilationStage::CodegenNoStubs {
                 target_harnesses: no_stubs,
-                next_harnesses: with_stubs,
+                next_harnesses: group_by_stubs(with_stubs, &all_harnesses),
                 all_harnesses,
             }
         } else {
@@ -266,7 +271,15 @@ impl KaniCompiler {
     fn prepare_codegen(&mut self) -> Compilation {
         debug!(stage=?self.stage, "prepare_codegen");
         match &self.stage {
-            CompilationStage::CodegenNoStubs { target_harnesses, all_harnesses, .. } => {
+            CompilationStage::CodegenNoStubs { target_harnesses, all_harnesses, .. }
+            | CompilationStage::CodegenWithStubs { target_harnesses, all_harnesses, .. } => {
+                debug!(
+                    harnesses=?target_harnesses
+                        .iter()
+                        .map(|h| &all_harnesses[h].metadata.pretty_name)
+                        .collect::<Vec<_>>(),
+                        "prepare_codegen"
+                );
                 let queries = &mut (*self.queries.lock().unwrap());
                 queries.harnesses_info = target_harnesses
                     .iter()
@@ -276,14 +289,6 @@ impl KaniCompiler {
                     .collect();
                 Compilation::Continue
             }
-            CompilationStage::CodegenWithStubs { target_harness, all_harnesses, .. } => {
-                let queries = &mut (*self.queries.lock().unwrap());
-                queries.harnesses_info = HashMap::from([(
-                    *target_harness,
-                    all_harnesses[&target_harness].metadata.goto_file.clone().unwrap(),
-                )]);
-                Compilation::Continue
-            }
             CompilationStage::Init | CompilationStage::Done => unreachable!(),
         }
     }
@@ -314,6 +319,18 @@ impl KaniCompiler {
     }
 }
 
+/// Group the harnesses by their stubs.
+fn group_by_stubs(
+    harnesses: Vec<HarnessId>,
+    all_harnesses: &HashMap<HarnessId, HarnessInfo>,
+) -> Vec<Vec<HarnessId>> {
+    let mut per_stubs: BTreeMap<&Stubs, Vec<HarnessId>> = BTreeMap::default();
+    for harness in harnesses {
+        per_stubs.entry(&all_harnesses[&harness].stub_map).or_default().push(harness)
+    }
+    per_stubs.into_values().collect()
+}
+
 /// Use default function implementations.
 impl Callbacks for KaniCompiler {
     /// Configure the [KaniCompiler] `self` object during the [CompilationStage::Init].
@@ -326,7 +343,6 @@ impl Callbacks for KaniCompiler {
                 &matches,
                 matches!(config.opts.error_format, ErrorOutputType::Json { .. }),
             );
-
             // Configure queries.
             let queries = &mut (*self.queries.lock().unwrap());
             queries.emit_vtable_restrictions = matches.get_flag(parser::RESTRICT_FN_PTRS);
@@ -364,3 +380,75 @@ impl Callbacks for KaniCompiler {
         self.prepare_codegen()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{HarnessInfo, Stubs};
+    use crate::kani_compiler::{group_by_stubs, HarnessId};
+    use kani_metadata::{HarnessAttributes, HarnessMetadata};
+    use rustc_data_structures::fingerprint::Fingerprint;
+    use rustc_hir::definitions::DefPathHash;
+    use std::collections::HashMap;
+
+    fn mock_next_id() -> HarnessId {
+        static mut COUNTER: u64 = 0;
+        unsafe { COUNTER += 1 };
+        let id = unsafe { COUNTER };
+        DefPathHash(Fingerprint::new(id, 0))
+    }
+
+    fn mock_metadata() -> HarnessMetadata {
+        HarnessMetadata {
+            pretty_name: String::from("dummy"),
+            mangled_name: String::from("dummy"),
+            crate_name: String::from("dummy"),
+            original_file: String::from("dummy"),
+            original_start_line: 10,
+            original_end_line: 20,
+            goto_file: None,
+            attributes: HarnessAttributes::default(),
+        }
+    }
+
+    fn mock_info_with_stubs(stub_map: Stubs) -> HarnessInfo {
+        HarnessInfo { metadata: mock_metadata(), stub_map }
+    }
+
+    #[test]
+    fn test_group_by_stubs_works() {
+        // Set up the inputs
+        let harness_1 = mock_next_id();
+        let harness_2 = mock_next_id();
+        let harness_3 = mock_next_id();
+        let harnesses = vec![harness_1, harness_2, harness_3];
+
+        let stub_1 = (mock_next_id(), mock_next_id());
+        let stub_2 = (mock_next_id(), mock_next_id());
+        let stub_3 = (mock_next_id(), mock_next_id());
+        let stub_4 = (stub_3.0, mock_next_id());
+
+        let set_1 = Stubs::from([stub_1, stub_2, stub_3]);
+        let set_2 = Stubs::from([stub_1, stub_2, stub_4]);
+        let set_3 = Stubs::from([stub_1, stub_3, stub_2]);
+        assert_eq!(set_1, set_3);
+        assert_ne!(set_1, set_2);
+
+        let harnesses_info = HashMap::from([
+            (harness_1, mock_info_with_stubs(set_1)),
+            (harness_2, mock_info_with_stubs(set_2)),
+            (harness_3, mock_info_with_stubs(set_3)),
+        ]);
+        assert_eq!(harnesses_info.len(), 3);
+
+        // Run the function under test.
+        let grouped = group_by_stubs(harnesses, &harnesses_info);
+
+        // Verify output.
+        assert_eq!(grouped.len(), 2);
+        assert!(
+            grouped.contains(&vec![harness_1, harness_3])
+                || grouped.contains(&vec![harness_3, harness_1])
+        );
+        assert!(grouped.contains(&vec![harness_2]));
+    }
+}

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

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 //! This file contains code for extracting stubbing-related attributes.
 
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 
 use kani_metadata::Stub;
 use rustc_hir::def_id::{DefId, LocalDefId};
@@ -42,7 +42,7 @@ pub fn update_stub_mapping(
     tcx: TyCtxt,
     harness: LocalDefId,
     stub: &Stub,
-    stub_pairs: &mut HashMap<DefPathHash, DefPathHash>,
+    stub_pairs: &mut BTreeMap<DefPathHash, DefPathHash>,
 ) {
     if let Some((orig_id, stub_id)) = stub_def_ids(tcx, harness, stub) {
         let orig_hash = tcx.def_path_hash(orig_id);

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

@@ -5,7 +5,7 @@
 mod annotations;
 mod transform;
 
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 
 use kani_metadata::HarnessMetadata;
 use rustc_hir::def_id::DefId;
@@ -20,9 +20,9 @@ pub fn harness_stub_map(
     tcx: TyCtxt,
     harness: DefId,
     metadata: &HarnessMetadata,
-) -> HashMap<DefPathHash, DefPathHash> {
+) -> BTreeMap<DefPathHash, DefPathHash> {
     let attrs = &metadata.attributes;
-    let mut stub_pairs = HashMap::default();
+    let mut stub_pairs = BTreeMap::default();
     for stubs in &attrs.stubs {
         update_stub_mapping(tcx, harness.expect_local(), stubs, &mut stub_pairs);
     }

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

@@ -6,7 +6,7 @@
 //! body of its stub, if appropriate. The stub mapping it uses is set via rustc
 //! arguments.
 
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 
 use lazy_static::lazy_static;
 use regex::Regex;
@@ -115,7 +115,7 @@ fn check_compatibility<'a, 'tcx>(
 const RUSTC_ARG_PREFIX: &str = "kani_stubs=";
 
 /// Serializes the stub mapping into a rustc argument.
-pub fn mk_rustc_arg(stub_mapping: &HashMap<DefPathHash, DefPathHash>) -> String {
+pub fn mk_rustc_arg(stub_mapping: &BTreeMap<DefPathHash, DefPathHash>) -> String {
     // Serialize each `DefPathHash` as a pair of `u64`s, and the whole mapping
     // as an association list.
     let mut pairs = Vec::new();

tests/script-based-pre/stubbing_compiler_sessions/check_compiler_sessions.expected

@@ -0,0 +1,3 @@
+Complete - 6 successfully verified harnesses, 0 failures, 6 total.
+
+Rust compiler sessions: 4

tests/script-based-pre/stubbing_compiler_sessions/check_compiler_sessions.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+#
+# Checks that the Kani compiler can encode harnesses with the same set of stubs
+# in one rustc session.
+
+set +e
+
+log_file=output.log
+
+kani stubbing.rs --enable-unstable --enable-stubbing --verbose >& ${log_file}
+
+echo "------- Raw output ---------"
+cat $log_file
+echo "----------------------------"
+
+# We print the reachability analysis results once for each session.
+# This is the only reliable way to get the number of sessions from the compiler.
+# The other option would be to use debug comments.
+# Ideally, the compiler should only print one set of statistics at the end of its run.
+# In that case, we should include number of sessions to those stats.
+runs=$(grep -c "Reachability Analysis Result" ${log_file})
+echo "Rust compiler sessions: ${runs}"
+
+# Cleanup
+rm ${log_file}

tests/script-based-pre/stubbing_compiler_sessions/config.yml

@@ -0,0 +1,4 @@
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+script: check_compiler_sessions.sh
+expected: check_compiler_sessions.expected

tests/script-based-pre/stubbing_compiler_sessions/stubbing.rs

@@ -0,0 +1,64 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// Check that Kani handles different sets of stubbing correctly.
+// I.e., not correctly replacing the stubs will cause a harness to fail.
+
+fn identity(i: i8) -> i8 {
+    i
+}
+
+fn decrement(i: i8) -> i8 {
+    i.wrapping_sub(1)
+}
+
+fn increment(i: i8) -> i8 {
+    i.wrapping_add(1)
+}
+
+#[kani::proof]
+fn check_identity() {
+    let n = kani::any();
+    assert_eq!(identity(n), n);
+}
+
+#[kani::proof]
+fn check_decrement() {
+    let n = kani::any();
+    kani::assume(n > i8::MIN);
+    assert_eq!(decrement(n), n - 1);
+}
+
+#[kani::proof]
+#[kani::stub(decrement, increment)]
+fn check_decrement_is_increment() {
+    let n = kani::any();
+    kani::assume(n < i8::MAX);
+    assert_eq!(decrement(n), n + 1);
+}
+
+#[kani::proof]
+#[kani::stub(increment, identity)]
+#[kani::stub(decrement, identity)]
+fn check_all_identity() {
+    let n = kani::any();
+    assert_eq!(decrement(n), increment(n));
+}
+
+#[kani::proof]
+#[kani::stub(decrement, identity)]
+#[kani::stub(increment, identity)]
+fn check_all_identity_2() {
+    let n = kani::any();
+    assert_eq!(decrement(n), n);
+    assert_eq!(increment(n), n);
+}
+
+#[kani::proof]
+#[kani::stub(decrement, increment)]
+#[kani::stub(increment, identity)]
+fn check_indirect_all_identity() {
+    let n = kani::any();
+    assert_eq!(decrement(n), n);
+    assert_eq!(increment(n), n);
+}