Moving all the contract communication logic into the compiler

kani-compiler/src/kani_compiler.rs

@@ -30,6 +30,7 @@ use clap::Parser;
 use kani_metadata::{ArtifactType, HarnessMetadata, KaniMetadata};
 use rustc_codegen_ssa::traits::CodegenBackend;
 use rustc_driver::{Callbacks, Compilation, RunCompiler};
+use rustc_hash::FxHashMap;
 use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_hir::definitions::DefPathHash;
 use rustc_interface::Config;
@@ -187,17 +188,26 @@ impl KaniCompiler {
             // Because this modifies `self.stage` we need to run this before
             // borrowing `&self.stage` immutably
             if let CompilationStage::Done { metadata: Some((metadata, _)), .. } = &mut self.stage {
-                let mut qdb = self.queries.lock().unwrap();
+                let mut contracts = self
+                    .queries
+                    .lock()
+                    .unwrap()
+                    .assigns_contracts()
+                    .map(|(k, v)| (*k, v.clone()))
+                    .collect::<FxHashMap<_, _>>();
                 for harness in
                     metadata.proof_harnesses.iter_mut().chain(metadata.test_harnesses.iter_mut())
                 {
                     if let Some(modifies_contract) =
-                        qdb.assigns_contract_for((&harness.mangled_name).intern())
+                        contracts.remove(&(&harness.mangled_name).intern())
                     {
                         harness.contract = modifies_contract.into();
                     }
                 }
-                qdb.assert_assigns_contracts_retrieved();
+                assert!(
+                    contracts.is_empty(),
+                    "Invariant broken: not all contracts have been handled."
+                )
             }
             match &self.stage {
                 CompilationStage::Init => self.run_compilation_session(&orig_args)?,

kani-compiler/src/kani_queries/mod.rs

@@ -59,26 +59,10 @@ impl QueryDb {
         )
     }
 
-    /// Lookup if a CBMC-level `assigns` contract was registered for this
-    /// harness with [`Self::add_assigns_contract`].
-    ///
-    /// This removes the contract from the registry and is intended to be used in
-    /// conjunction with [`Self::assert_assigns_contracts_retrieved`] to uphold
-    /// the invariant that each contract has been handled and only handled once.
-    pub fn assigns_contract_for(
-        &mut self,
-        harness_name: InternedString,
-    ) -> Option<AssignsContract> {
-        self.modifies_contracts.remove(&harness_name)
-    }
-
-    /// Assert that the contract registry is empty. See [`Self::assigns_contract_for`]
-    pub fn assert_assigns_contracts_retrieved(&self) {
-        assert!(
-            self.modifies_contracts.is_empty(),
-            "Invariant broken: The modifies contracts for {} have not been retrieved",
-            PrintList(self.modifies_contracts.keys())
-        )
+    /// Lookup all CBMC-level `assigns` contract were registered with
+    /// [`Self::add_assigns_contract`].
+    pub fn assigns_contracts(&self) -> impl Iterator<Item = (&InternedString, &AssignsContract)> {
+        self.modifies_contracts.iter()
     }
 }