Clone a template BodyTransformer to avoid re-initialization (#4259)

Currently, we have to create a new `BodyTransformer` for each harness
individually. However, this takes a lot of time (as we have to
constantly re-validate all our Kani functions and reinitialize
everything based on the `QueryDb`) and the transformer's options are all
shared between harnesses of the same codegen unit anyway.

This PR makes the `BodyTransformer` struct `Clone`-able, allowing us to
initialize a single 'template' transformer for each codegen unit and
then cheaply clone the template for each harness within the unit. Based
on testing, the clone takes just ~3µs when using the default # of
transformation passes, which is much faster than initialization.

### Results
This change made a noticeable **4.8% reduction in the end to end compile
time of the standard library** (at std commit
[177d0fd](model-checking/verify-rust-std@177d0fd)
& assuming #4257 is merged into Kani).

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

Author: AlexanderPortland

kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs

@@ -90,6 +90,7 @@ impl GotocCodegenBackend {
         symtab_goto: &Path,
         machine_model: &MachineModel,
         check_contract: Option<InternalDefId>,
+        mut global_passes: GlobalPasses,
         mut transformer: BodyTransformation,
         thread_pool: &ThreadPool,
     ) -> (MinimalGotocCtx, Vec<MonoItem>, Option<AssignsContract>) {
@@ -119,7 +120,6 @@ impl GotocCodegenBackend {
             .collect();
 
         // Apply all transformation passes, including global passes.
-        let mut global_passes = GlobalPasses::new(&self.queries.lock().unwrap(), tcx);
         let any_pass_modified = global_passes.run_global_passes(
             &mut transformer,
             tcx,
@@ -393,12 +393,17 @@ impl CodegenBackend for GotocCodegenBackend {
                     let num_harnesses: usize = units.iter().map(|unit| unit.harnesses.len()).sum();
                     export_thread_pool.add_workers(Self::thread_pool_size(Some(num_harnesses)));
 
+                    let template_passes = GlobalPasses::new(&queries, tcx);
+
                     // Cross-crate collecting of all items that are reachable from the crate harnesses.
                     for unit in units.iter() {
                         // We reset the body cache for now because each codegen unit has different
                         // configurations that affect how we transform the instance body.
+
+                        // Generate an empty 'template' transformer once per codegen unit and then clone for each harness within.
+                        // (They all share the same options.)
+                        let template_transformer = BodyTransformation::new(&queries, tcx, unit);
                         for harness in &unit.harnesses {
-                            let transformer = BodyTransformation::new(&queries, tcx, unit);
                             let model_path = units.harness_model_path(*harness).unwrap();
                             let is_automatic_harness = units.is_automatic_harness(harness);
                             let contract_metadata =
@@ -410,7 +415,8 @@ impl CodegenBackend for GotocCodegenBackend {
                                 &results.machine_model,
                                 contract_metadata
                                     .map(|def| rustc_internal::internal(tcx, def.def_id())),
-                                transformer,
+                                template_passes.clone(),
+                                template_transformer.clone_empty(),
                                 &export_thread_pool,
                             );
                             if min_gcx.has_loop_contracts {
@@ -450,6 +456,7 @@ impl CodegenBackend for GotocCodegenBackend {
                         &model_path,
                         &results.machine_model,
                         Default::default(),
+                        GlobalPasses::new(&queries, tcx),
                         transformer,
                         &export_thread_pool,
                     );

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

@@ -30,7 +30,7 @@ use tracing::debug;
 
 /// Generate `T::any()` implementations for `T`s that do not implement Arbitrary in source code.
 /// Currently limited to structs and enums.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct AutomaticArbitraryPass {
     /// The FnDef of KaniModel::Any
     kani_any: FnDef,
@@ -282,7 +282,7 @@ impl AutomaticArbitraryPass {
     }
 }
 /// Transform the dummy body of an automatic_harness Kani intrinsic to be a proof harness for a given function.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct AutomaticHarnessPass {
     kani_any: FnDef,
     init_contracts_hook: Instance,

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

@@ -31,7 +31,7 @@ use rustc_session::config::OutputType;
 mod initial_target_visitor;
 mod instrumentation_visitor;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct DelayedUbPass {
     pub safety_check_type: CheckType,
     pub unsupported_check_type: CheckType,

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

@@ -28,7 +28,7 @@ mod uninit_visitor;
 
 /// Top-level pass that instruments the code with checks for uninitialized memory access through raw
 /// pointers.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct UninitPass {
     pub safety_check_type: CheckType,
     pub unsupported_check_type: CheckType,

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

@@ -38,7 +38,7 @@ use strum_macros::AsRefStr;
 use tracing::{debug, trace};
 
 /// Instrument the code with checks for invalid values.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct ValidValuePass {
     pub safety_check_type: CheckType,
     pub unsupported_check_type: CheckType,

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

@@ -31,7 +31,7 @@ use tracing::{debug, trace};
 /// depending on what the type of the input it
 ///
 /// any_modifies is replaced with any
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct AnyModifiesPass {
     kani_any: Option<FnDef>,
     kani_any_modifies: Option<FnDef>,
@@ -265,7 +265,7 @@ impl AnyModifiesPass {
 ///    - Replace the non-used generated closures body with unreachable.
 /// 3. Replace the body of `kani_register_contract` by `kani::internal::run_contract_fn` to
 ///    invoke the closure.
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct FunctionWithContractPass {
     /// Function that is being checked, if any.
     check_fn: Option<FnDef>,

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

@@ -17,7 +17,7 @@ use std::io::Write;
 use super::BodyTransformation;
 
 /// Dump all MIR bodies.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct DumpMirPass {
     enabled: bool,
 }

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

@@ -37,7 +37,7 @@ use std::str::FromStr;
 use tracing::{debug, trace};
 
 /// Generate the body for a few Kani intrinsics.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct IntrinsicGeneratorPass {
     unsupported_check_type: CheckType,
     /// Used to cache FnDef lookups for models and Kani intrinsics.

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

@@ -23,7 +23,7 @@ use rustc_span::Symbol;
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::fmt::Debug;
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct LoopContractPass {
     /// Cache KaniRunContract function used to implement contracts.
     run_contract_fn: Option<FnDef>,

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

@@ -21,6 +21,7 @@ use crate::kani_middle::reachability::CallGraph;
 use crate::kani_middle::transform::body::CheckType;
 use crate::kani_middle::transform::check_uninit::{DelayedUbPass, UninitPass};
 use crate::kani_middle::transform::check_values::ValidValuePass;
+use crate::kani_middle::transform::clone::{ClonableGlobalPass, ClonableTransformPass};
 use crate::kani_middle::transform::contracts::{AnyModifiesPass, FunctionWithContractPass};
 use crate::kani_middle::transform::kani_intrinsics::IntrinsicGeneratorPass;
 use crate::kani_middle::transform::loop_contracts::LoopContractPass;
@@ -58,9 +59,9 @@ mod stubs;
 pub struct BodyTransformation {
     /// The passes that may change the function body according to harness configuration.
     /// The stubbing passes should be applied before so user stubs take precedence.
-    stub_passes: Vec<Box<dyn TransformPass>>,
+    stub_passes: Vec<Box<dyn ClonableTransformPass>>,
     /// The passes that may add safety checks to the function body.
-    inst_passes: Vec<Box<dyn TransformPass>>,
+    inst_passes: Vec<Box<dyn ClonableTransformPass>>,
     /// Cache transformation results.
     cache: HashMap<Instance, TransformationResult>,
 }
@@ -86,7 +87,7 @@ impl BodyTransformation {
         transformer.add_pass(
             queries,
             ValidValuePass {
-                safety_check_type: safety_check_type.clone(),
+                safety_check_type,
                 unsupported_check_type: unsupported_check_type.clone(),
             },
         );
@@ -139,7 +140,18 @@ impl BodyTransformation {
         }
     }
 
-    fn add_pass<P: TransformPass + 'static>(&mut self, query_db: &QueryDb, pass: P) {
+    /// Clone an empty [BodyTransformation] for use within the same [CodegenUnit] and [TyCtxt] that were
+    /// used to create it. Will panic if the transformer has already been queried with `.body()`.
+    pub fn clone_empty(&self) -> Self {
+        debug_assert!(self.cache.is_empty());
+        BodyTransformation {
+            stub_passes: self.stub_passes.to_vec(),
+            inst_passes: self.inst_passes.to_vec(),
+            cache: HashMap::new(),
+        }
+    }
+
+    fn add_pass<P: ClonableTransformPass + 'static>(&mut self, query_db: &QueryDb, pass: P) {
         if pass.is_enabled(query_db) {
             match P::transformation_type() {
                 TransformationType::Instrumentation => self.inst_passes.push(Box::new(pass)),
@@ -199,10 +211,11 @@ enum TransformationResult {
     NotModified,
 }
 
+#[derive(Clone)]
 pub struct GlobalPasses {
     /// The passes that operate on the whole codegen unit, they run after all previous passes are
     /// done.
-    global_passes: Vec<Box<dyn GlobalPass>>,
+    global_passes: Vec<Box<dyn ClonableGlobalPass>>,
 }
 
 impl GlobalPasses {
@@ -220,7 +233,7 @@ impl GlobalPasses {
         global_passes
     }
 
-    fn add_global_pass<P: GlobalPass + 'static>(&mut self, query_db: &QueryDb, pass: P) {
+    fn add_global_pass<P: ClonableGlobalPass + 'static>(&mut self, query_db: &QueryDb, pass: P) {
         if pass.is_enabled(query_db) {
             self.global_passes.push(Box::new(pass))
         }
@@ -249,3 +262,57 @@ impl GlobalPasses {
         modified
     }
 }
+
+mod clone {
+    //! This is all just machinery to implement `Clone` for a `Box<dyn TransformPass + Clone>`.
+    //!
+    //! To avoid circular reasoning, we use two traits that can each clone into a dyn of the other, and since
+    //! we set both up to have blanket implementations for all `T: TransformPass + Clone`, the compiler pieces them together properly
+    //! and we can implement `Clone` directly using the pair!
+
+    /// Builds a new dyn compatible wrapper trait that's essentially equivalent to extending
+    /// `$extends` with `Clone`. Requires an ident for an intermediate trait for avoiding cycles
+    /// in the implementation.
+    macro_rules! implement_clone {
+        ($new_trait_name: ident, $intermediate_trait_name: ident, $extends: path) => {
+            #[allow(private_interfaces)]
+            pub(crate) trait $new_trait_name: $extends {
+                fn clone_there(&self) -> Box<dyn $intermediate_trait_name>;
+            }
+
+            impl Clone for Box<dyn $new_trait_name> {
+                fn clone(&self) -> Self {
+                    self.clone_there().clone_back()
+                }
+            }
+
+            #[allow(private_interfaces)]
+            impl<T: $extends + Clone + 'static> $new_trait_name for T {
+                fn clone_there(&self) -> Box<dyn $intermediate_trait_name> {
+                    Box::new(self.clone())
+                }
+            }
+
+            trait $intermediate_trait_name {
+                fn clone_back(&self) -> Box<dyn $new_trait_name>;
+            }
+
+            impl<T: $extends + Clone + 'static> $intermediate_trait_name for T {
+                fn clone_back(&self) -> Box<dyn $new_trait_name> {
+                    Box::new(self.clone())
+                }
+            }
+        };
+    }
+
+    implement_clone!(
+        ClonableTransformPass,
+        ClonableTransformPassMid,
+        crate::kani_middle::transform::TransformPass
+    );
+    implement_clone!(
+        ClonableGlobalPass,
+        ClonableGlobalPassMid,
+        crate::kani_middle::transform::GlobalPass
+    );
+}