From d0c0e31a3d0c1ca09247e596b53f223569eab182 Mon Sep 17 00:00:00 2001
From: "Felipe R. Monteiro" <felisous@amazon.com>
Date: Fri, 4 Aug 2023 19:46:54 +0000
Subject: [PATCH] Only replace call to foreign functions

Signed-off-by: Felipe R. Monteiro <felisous@amazon.com>
---
 kani-compiler/src/kani_middle/provide.rs      |  2 +-
 .../src/kani_middle/stubbing/transform.rs     | 40 +++++++++++++------
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/kani-compiler/src/kani_middle/provide.rs b/kani-compiler/src/kani_middle/provide.rs
index 8d119b007514..c07cc1e477fe 100644
--- a/kani-compiler/src/kani_middle/provide.rs
+++ b/kani-compiler/src/kani_middle/provide.rs
@@ -56,7 +56,7 @@ fn run_kani_mir_passes<'tcx>(
 ) -> &'tcx Body<'tcx> {
     tracing::debug!(?def_id, "Run Kani transformation passes");
     let mut transformed_body = stubbing::transform(tcx, def_id, body);
-    stubbing::transform_extern_functions(tcx, &mut transformed_body);
+    stubbing::transform_foreign_functions(tcx, &mut transformed_body);
     tcx.arena.alloc(transformed_body)
 }
 
diff --git a/kani-compiler/src/kani_middle/stubbing/transform.rs b/kani-compiler/src/kani_middle/stubbing/transform.rs
index 1e5154f03b37..6a3ab74af978 100644
--- a/kani-compiler/src/kani_middle/stubbing/transform.rs
+++ b/kani-compiler/src/kani_middle/stubbing/transform.rs
@@ -36,12 +36,29 @@ pub fn transform<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, old_body: &'tcx Body<'t
     old_body.clone()
 }
 
-struct ExternFunctionTransformer<'tcx> {
+/// Tranverse `body` searching for calls to foreing functions and, whevever there is
+/// a stub available, replace the call to the foreign function with a call
+/// to its correspondent stub. This happens as a separate step because there is no
+/// body available to foreign functions at this stage.
+pub fn transform_foreign_functions<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    let mut visitor = if let Some(stub_map) = get_stub_mapping(tcx) {
+        ForeignFunctionTransformer { tcx, body: body.clone(), stub_map }
+    } else {
+        ForeignFunctionTransformer { tcx, body: body.clone(), stub_map: HashMap::default() }
+    };
+    visitor.visit_body(body);
+}
+
+struct ForeignFunctionTransformer<'tcx> {
+    /// The compiler context.
     tcx: TyCtxt<'tcx>,
+    /// Body of the callee function. Kani searches here for foreign functions.
     body: Body<'tcx>,
+    /// Map of functions/methods to their correspondent stubs.
+    stub_map: HashMap<DefId, DefId>,
 }
 
-impl<'tcx> MutVisitor<'tcx> for ExternFunctionTransformer<'tcx> {
+impl<'tcx> MutVisitor<'tcx> for ForeignFunctionTransformer<'tcx> {
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
@@ -49,22 +66,19 @@ impl<'tcx> MutVisitor<'tcx> for ExternFunctionTransformer<'tcx> {
     fn visit_operand(&mut self, operand: &mut Operand<'tcx>, _location: Location) {
         let func_ty = operand.ty(&self.body, self.tcx);
         if let ty::FnDef(reachable_function, generics) = *func_ty.kind() {
-            if let Some(stub) = get_stub(self.tcx, reachable_function) {
-                let Operand::Constant(fn_def) = operand else { unreachable!() };
-                fn_def.literal = ConstantKind::from_value(
-                    ConstValue::ZeroSized,
-                    self.tcx.type_of(stub).subst(self.tcx, generics),
-                );
+            if self.tcx.is_foreign_item(reachable_function) {
+                if let Some(stub) = self.stub_map.get(&reachable_function) {
+                    let Operand::Constant(fn_def) = operand else { unreachable!() };
+                    fn_def.literal = ConstantKind::from_value(
+                        ConstValue::ZeroSized,
+                        self.tcx.type_of(stub).subst(self.tcx, generics),
+                    );
+                }
             }
         }
     }
 }
 
-pub fn transform_extern_functions<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    let mut visitor = ExternFunctionTransformer { tcx, body: body.clone() };
-    visitor.visit_body(body);
-}
-
 /// Checks whether the stub is compatible with the original function/method: do
 /// the arities and types (of the parameters and return values) match up? This
 /// does **NOT** check whether the type variables are constrained to implement