From e1777710b1b40fb6e4804e9a81253a9d30228bc4 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 6 Apr 2023 20:44:12 +0200
Subject: [PATCH] Don't override declare_*_in_func in cranelift-jit (#6169)

Instead remove the colocated flag for hotplug mode in define_function.
This prevents issues if declare_*_in_func wasn't used due to eg the
function being from a previously serialized module and now deserialized
into JITModule.
---
 cranelift/jit/src/backend.rs | 52 +++++++++++-------------------------
 1 file changed, 15 insertions(+), 37 deletions(-)

diff --git a/cranelift/jit/src/backend.rs b/cranelift/jit/src/backend.rs
index 79e5db86af..7defa0de73 100644
--- a/cranelift/jit/src/backend.rs
+++ b/cranelift/jit/src/backend.rs
@@ -642,43 +642,6 @@ impl Module for JITModule {
         Ok(id)
     }
 
-    /// Use this when you're building the IR of a function to reference a function.
-    ///
-    /// TODO: Coalesce redundant decls and signatures.
-    /// TODO: Look into ways to reduce the risk of using a FuncRef in the wrong function.
-    fn declare_func_in_func(&mut self, func: FuncId, in_func: &mut ir::Function) -> ir::FuncRef {
-        let decl = self.declarations.get_function_decl(func);
-        let signature = in_func.import_signature(decl.signature.clone());
-        let colocated = !self.hotswap_enabled && decl.linkage.is_final();
-        let user_name_ref = in_func.declare_imported_user_function(ir::UserExternalName {
-            namespace: 0,
-            index: func.as_u32(),
-        });
-        in_func.import_function(ir::ExtFuncData {
-            name: ir::ExternalName::user(user_name_ref),
-            signature,
-            colocated,
-        })
-    }
-
-    /// Use this when you're building the IR of a function to reference a data object.
-    ///
-    /// TODO: Same as above.
-    fn declare_data_in_func(&self, data: DataId, func: &mut ir::Function) -> ir::GlobalValue {
-        let decl = self.declarations.get_data_decl(data);
-        let colocated = !self.hotswap_enabled && decl.linkage.is_final();
-        let user_name_ref = func.declare_imported_user_function(ir::UserExternalName {
-            namespace: 1,
-            index: data.as_u32(),
-        });
-        func.create_global_value(ir::GlobalValueData::Symbol {
-            name: ir::ExternalName::user(user_name_ref),
-            offset: ir::immediates::Imm64::new(0),
-            colocated,
-            tls: decl.tls,
-        })
-    }
-
     fn define_function_with_control_plane(
         &mut self,
         id: FuncId,
@@ -695,6 +658,21 @@ impl Module for JITModule {
             return Err(ModuleError::DuplicateDefinition(decl.name.to_owned()));
         }
 
+        if self.hotswap_enabled {
+            // Disable colocated if hotswapping is enabled to avoid a PLT indirection in case of
+            // calls and to allow data objects to be hotswapped in the future.
+            for func in ctx.func.dfg.ext_funcs.values_mut() {
+                func.colocated = false;
+            }
+
+            for gv in ctx.func.global_values.values_mut() {
+                match gv {
+                    ir::GlobalValueData::Symbol { colocated, .. } => *colocated = false,
+                    _ => {}
+                }
+            }
+        }
+
         // work around borrow-checker to allow reuse of ctx below
         let res = ctx.compile(self.isa(), ctrl_plane)?;
         let alignment = res.alignment as u64;