From bbc2afb7a6197c5d4d42a5ec0b80f3c4032bb713 Mon Sep 17 00:00:00 2001
From: bjorn3 <bjorn3@users.noreply.github.com>
Date: Thu, 12 Nov 2020 19:29:05 +0100
Subject: [PATCH] Add prepare_for_function_redefine

---
 Cargo.lock                         |  1 +
 cranelift/module/src/module.rs     | 94 ++++++++++++++++++++++++++++++
 cranelift/simplejit/Cargo.toml     |  1 +
 cranelift/simplejit/src/backend.rs | 19 ++++++
 4 files changed, 115 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 7e22a7d85aaa..36dc8c1618a1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -477,6 +477,7 @@ dependencies = [
 name = "cranelift-simplejit"
 version = "0.68.0"
 dependencies = [
+ "anyhow",
  "cranelift",
  "cranelift-codegen",
  "cranelift-entity",
diff --git a/cranelift/module/src/module.rs b/cranelift/module/src/module.rs
index eb4e668b98bc..735704ffa17d 100644
--- a/cranelift/module/src/module.rs
+++ b/cranelift/module/src/module.rs
@@ -490,3 +490,97 @@ pub trait Module {
     /// Define a data object, producing the data contents from the given `DataContext`.
     fn define_data(&mut self, data: DataId, data_ctx: &DataContext) -> ModuleResult<()>;
 }
+
+impl<M: Module> Module for &mut M {
+    fn isa(&self) -> &dyn isa::TargetIsa {
+        (**self).isa()
+    }
+
+    fn declarations(&self) -> &ModuleDeclarations {
+        (**self).declarations()
+    }
+
+    fn get_name(&self, name: &str) -> Option<FuncOrDataId> {
+        (**self).get_name(name)
+    }
+
+    fn target_config(&self) -> isa::TargetFrontendConfig {
+        (**self).target_config()
+    }
+
+    fn make_context(&self) -> Context {
+        (**self).make_context()
+    }
+
+    fn clear_context(&self, ctx: &mut Context) {
+        (**self).clear_context(ctx)
+    }
+
+    fn make_signature(&self) -> ir::Signature {
+        (**self).make_signature()
+    }
+
+    fn clear_signature(&self, sig: &mut ir::Signature) {
+        (**self).clear_signature(sig)
+    }
+
+    fn declare_function(
+        &mut self,
+        name: &str,
+        linkage: Linkage,
+        signature: &ir::Signature,
+    ) -> ModuleResult<FuncId> {
+        (**self).declare_function(name, linkage, signature)
+    }
+
+    fn declare_data(
+        &mut self,
+        name: &str,
+        linkage: Linkage,
+        writable: bool,
+        tls: bool,
+    ) -> ModuleResult<DataId> {
+        (**self).declare_data(name, linkage, writable, tls)
+    }
+
+    fn declare_func_in_func(&self, func: FuncId, in_func: &mut ir::Function) -> ir::FuncRef {
+        (**self).declare_func_in_func(func, in_func)
+    }
+
+    fn declare_data_in_func(&self, data: DataId, func: &mut ir::Function) -> ir::GlobalValue {
+        (**self).declare_data_in_func(data, func)
+    }
+
+    fn declare_func_in_data(&self, func: FuncId, ctx: &mut DataContext) -> ir::FuncRef {
+        (**self).declare_func_in_data(func, ctx)
+    }
+
+    fn declare_data_in_data(&self, data: DataId, ctx: &mut DataContext) -> ir::GlobalValue {
+        (**self).declare_data_in_data(data, ctx)
+    }
+
+    fn define_function<TS>(
+        &mut self,
+        func: FuncId,
+        ctx: &mut Context,
+        trap_sink: &mut TS,
+    ) -> ModuleResult<ModuleCompiledFunction>
+    where
+        TS: binemit::TrapSink
+    {
+        (**self).define_function(func, ctx, trap_sink)
+    }
+
+    fn define_function_bytes(
+        &mut self,
+        func: FuncId,
+        bytes: &[u8],
+        relocs: &[RelocRecord],
+    ) -> ModuleResult<ModuleCompiledFunction> {
+        (**self).define_function_bytes(func, bytes, relocs)
+    }
+
+    fn define_data(&mut self, data: DataId, data_ctx: &DataContext) -> ModuleResult<()> {
+        (**self).define_data(data, data_ctx)
+    }
+}
diff --git a/cranelift/simplejit/Cargo.toml b/cranelift/simplejit/Cargo.toml
index 9c7bea2c3dd7..ea8165035a42 100644
--- a/cranelift/simplejit/Cargo.toml
+++ b/cranelift/simplejit/Cargo.toml
@@ -14,6 +14,7 @@ cranelift-module = { path = "../module", version = "0.68.0" }
 cranelift-native = { path = "../native", version = "0.68.0" }
 cranelift-codegen = { path = "../codegen", version = "0.68.0", default-features = false, features = ["std"] }
 cranelift-entity = { path = "../entity", version = "0.68.0" }
+anyhow = "1.0"
 region = "2.2.0"
 libc = { version = "0.2.42" }
 errno = "0.2.4"
diff --git a/cranelift/simplejit/src/backend.rs b/cranelift/simplejit/src/backend.rs
index bb7ddef6e1d0..6602a8bd0571 100644
--- a/cranelift/simplejit/src/backend.rs
+++ b/cranelift/simplejit/src/backend.rs
@@ -393,6 +393,25 @@ impl SimpleJITModule {
             data_objects_to_finalize: Vec::new(),
         }
     }
+
+    /// Allow a single future `define_function` on a previously defined function. This allows for
+    /// hot code swapping and lazy compilation of functions.
+    pub fn prepare_for_function_redefine(&mut self, func_id: FuncId) -> ModuleResult<()> {
+        let decl = self.declarations.get_function_decl(func_id);
+        if !decl.linkage.is_definable() {
+            return Err(ModuleError::InvalidImportDefinition(decl.name.clone()));
+        }
+
+        if self.compiled_functions[func_id].is_none() {
+            return Err(ModuleError::Backend(anyhow::anyhow!("Tried to redefine not yet defined function {}", decl.name)));
+        }
+
+        self.compiled_functions[func_id] = None;
+
+        // FIXME return some kind of handle that allows for deallocating the function
+
+        Ok(())
+    }
 }
 
 impl<'simple_jit_backend> Module for SimpleJITModule {