1 files changed, 115 insertions, 0 deletions
diff --git a/compiler/rustc_codegen_cranelift/src/unwind_module.rs b/compiler/rustc_codegen_cranelift/src/unwind_module.rs
new file mode 100644
index 00000000000..b950aaa29ce
--- /dev/null
+++ b/compiler/rustc_codegen_cranelift/src/unwind_module.rs
@@ -0,0 +1,115 @@
+use cranelift_codegen::control::ControlPlane;
+use cranelift_codegen::ir::{Function, Signature};
+use cranelift_codegen::isa::{TargetFrontendConfig, TargetIsa};
+use cranelift_codegen::{Context, FinalizedMachReloc};
+use cranelift_module::{
+    DataDescription, DataId, FuncId, FuncOrDataId, Linkage, Module, ModuleDeclarations,
+    ModuleResult,
+};
+use cranelift_object::{ObjectModule, ObjectProduct};
+
+use crate::UnwindContext;
+
+/// A wrapper around a [Module] which adds any defined function to the [UnwindContext].
+pub(crate) struct UnwindModule<T> {
+    pub(crate) module: T,
+    unwind_context: UnwindContext,
+}
+
+impl<T: Module> UnwindModule<T> {
+    pub(crate) fn new(module: T, pic_eh_frame: bool) -> Self {
+        let unwind_context = UnwindContext::new(module.isa(), pic_eh_frame);
+        UnwindModule { module, unwind_context }
+    }
+}
+
+impl UnwindModule<ObjectModule> {
+    pub(crate) fn finish(self) -> ObjectProduct {
+        let mut product = self.module.finish();
+        self.unwind_context.emit(&mut product);
+        product
+    }
+}
+
+#[cfg(feature = "jit")]
+impl UnwindModule<cranelift_jit::JITModule> {
+    pub(crate) fn finalize_definitions(&mut self) {
+        self.module.finalize_definitions().unwrap();
+        let prev_unwind_context = std::mem::replace(
+            &mut self.unwind_context,
+            UnwindContext::new(self.module.isa(), false),
+        );
+        unsafe { prev_unwind_context.register_jit(&self.module) };
+    }
+}
+
+impl<T: Module> Module for UnwindModule<T> {
+    fn isa(&self) -> &dyn TargetIsa {
+        self.module.isa()
+    }
+
+    fn declarations(&self) -> &ModuleDeclarations {
+        self.module.declarations()
+    }
+
+    fn get_name(&self, name: &str) -> Option<FuncOrDataId> {
+        self.module.get_name(name)
+    }
+
+    fn target_config(&self) -> TargetFrontendConfig {
+        self.module.target_config()
+    }
+
+    fn declare_function(
+        &mut self,
+        name: &str,
+        linkage: Linkage,
+        signature: &Signature,
+    ) -> ModuleResult<FuncId> {
+        self.module.declare_function(name, linkage, signature)
+    }
+
+    fn declare_anonymous_function(&mut self, signature: &Signature) -> ModuleResult<FuncId> {
+        self.module.declare_anonymous_function(signature)
+    }
+
+    fn declare_data(
+        &mut self,
+        name: &str,
+        linkage: Linkage,
+        writable: bool,
+        tls: bool,
+    ) -> ModuleResult<DataId> {
+        self.module.declare_data(name, linkage, writable, tls)
+    }
+
+    fn declare_anonymous_data(&mut self, writable: bool, tls: bool) -> ModuleResult<DataId> {
+        self.module.declare_anonymous_data(writable, tls)
+    }
+
+    fn define_function_with_control_plane(
+        &mut self,
+        func: FuncId,
+        ctx: &mut Context,
+        ctrl_plane: &mut ControlPlane,
+    ) -> ModuleResult<()> {
+        self.module.define_function_with_control_plane(func, ctx, ctrl_plane)?;
+        self.unwind_context.add_function(func, ctx, self.module.isa());
+        Ok(())
+    }
+
+    fn define_function_bytes(
+        &mut self,
+        _func_id: FuncId,
+        _func: &Function,
+        _alignment: u64,
+        _bytes: &[u8],
+        _relocs: &[FinalizedMachReloc],
+    ) -> ModuleResult<()> {
+        unimplemented!()
+    }
+
+    fn define_data(&mut self, data_id: DataId, data: &DataDescription) -> ModuleResult<()> {
+        self.module.define_data(data_id, data)
+    }
+}