From 3bf51f0fbce68867d3213f3e02ba839e281332ad Mon Sep 17 00:00:00 2001
From: Artem Agvanian <nartagva@amazon.com>
Date: Fri, 5 Jul 2024 15:22:48 -0700
Subject: [PATCH] Wrap memory initialization set-up in a separate function

---
 .../kani_middle/transform/check_uninit/mod.rs | 74 ++++++++++---------
 1 file changed, 40 insertions(+), 34 deletions(-)

diff --git a/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs b/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
index dc86e5ab8f48..5f4610076bf1 100644
--- a/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
+++ b/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
@@ -76,40 +76,7 @@ impl TransformPass for UninitPass {
 
         // Inject a call to set-up memory initialization state if the function is a harness.
         if is_harness(instance, tcx) {
-            // First statement or terminator in the harness.
-            let mut source = if !new_body.blocks()[0].statements.is_empty() {
-                SourceInstruction::Statement { idx: 0, bb: 0 }
-            } else {
-                SourceInstruction::Terminator { bb: 0 }
-            };
-
-            // Dummy return place.
-            let ret_place = Place {
-                local: new_body.new_local(
-                    Ty::new_tuple(&[]),
-                    source.span(new_body.blocks()),
-                    Mutability::Not,
-                ),
-                projection: vec![],
-            };
-
-            // Resolve the instance and inject a call to set-up the memory initialization state.
-            let memory_initialization_init = Instance::resolve(
-                get_mem_init_fn_def(
-                    tcx,
-                    "KaniInitializeMemoryInitializationState",
-                    &mut self.mem_init_fn_cache,
-                ),
-                &GenericArgs(vec![]),
-            )
-            .unwrap();
-            new_body.add_call(
-                &memory_initialization_init,
-                &mut source,
-                InsertPosition::Before,
-                vec![],
-                ret_place,
-            );
+            inject_memory_init_setup(&mut new_body, tcx, &mut self.mem_init_fn_cache);
         }
 
         // Set of basic block indices for which analyzing first statement should be skipped.
@@ -481,3 +448,42 @@ fn is_harness(instance: Instance, tcx: TyCtxt) -> bool {
         tcx.has_attrs_with_path(rustc_internal::internal(tcx, instance.def.def_id()), attr_path)
     })
 }
+
+/// Inject an initial call to set-up memory initialization tracking.
+fn inject_memory_init_setup(
+    new_body: &mut MutableBody,
+    tcx: TyCtxt,
+    mem_init_fn_cache: &mut HashMap<&'static str, FnDef>,
+) {
+    // First statement or terminator in the harness.
+    let mut source = if !new_body.blocks()[0].statements.is_empty() {
+        SourceInstruction::Statement { idx: 0, bb: 0 }
+    } else {
+        SourceInstruction::Terminator { bb: 0 }
+    };
+
+    // Dummy return place.
+    let ret_place = Place {
+        local: new_body.new_local(
+            Ty::new_tuple(&[]),
+            source.span(new_body.blocks()),
+            Mutability::Not,
+        ),
+        projection: vec![],
+    };
+
+    // Resolve the instance and inject a call to set-up the memory initialization state.
+    let memory_initialization_init = Instance::resolve(
+        get_mem_init_fn_def(tcx, "KaniInitializeMemoryInitializationState", mem_init_fn_cache),
+        &GenericArgs(vec![]),
+    )
+    .unwrap();
+
+    new_body.add_call(
+        &memory_initialization_init,
+        &mut source,
+        InsertPosition::Before,
+        vec![],
+        ret_place,
+    );
+}