Introduce UnwindModule wrapper around a Module and UnwindContext

This way all UnwindContext::add_function calls can be done automatically
in a single place.

src/allocator.rs

@@ -11,15 +11,10 @@
 use crate::prelude::*;
 
 /// Returns whether an allocator shim was created
-pub(crate) fn codegen(
+pub(crate) fn codegen(tcx: TyCtxt<'_>, module: &mut impl Module) -> bool {
-    tcx: TyCtxt<'_>,
-    module: &mut impl Module,
-    unwind_context: &mut UnwindContext,
-) -> bool {
     let Some(kind) = allocator_kind_for_codegen(tcx) else { return false };
     codegen_inner(
         module,
-        unwind_context,
         kind,
         tcx.alloc_error_handler_kind(()).unwrap(),
         tcx.sess.opts.unstable_opts.oom,
@@ -29,7 +24,6 @@ pub(crate) fn codegen(
 
 fn codegen_inner(
     module: &mut impl Module,
-    unwind_context: &mut UnwindContext,
     kind: AllocatorKind,
     alloc_error_handler_kind: AllocatorKind,
     oom_strategy: OomStrategy,
@@ -67,7 +61,6 @@ fn codegen_inner(
             };
             crate::common::create_wrapper_function(
                 module,
-                unwind_context,
                 sig,
                 &global_fn_name(method.name),
                 &default_fn_name(method.name),
@@ -82,7 +75,6 @@ fn codegen_inner(
     };
     crate::common::create_wrapper_function(
         module,
-        unwind_context,
         sig,
         "__rust_alloc_error_handler",
         &alloc_error_handler_name(alloc_error_handler_kind),

src/base.rs

@@ -249,9 +249,7 @@ pub(crate) fn compile_fn(
     }
 
     // Define debuginfo for function
-    let isa = module.isa();
     let debug_context = &mut cx.debug_context;
-    let unwind_context = &mut cx.unwind_context;
     cx.profiler.generic_activity("generate debug info").run(|| {
         if let Some(debug_context) = debug_context {
             codegened_func.func_debug_cx.unwrap().finalize(
@@ -260,7 +258,6 @@ pub(crate) fn compile_fn(
                 context,
             );
         }
-        unwind_context.add_function(codegened_func.func_id, &context, isa);
     });
 }
 

src/common.rs

@@ -247,7 +247,6 @@ pub(crate) fn type_sign(ty: Ty<'_>) -> bool {
 
 pub(crate) fn create_wrapper_function(
     module: &mut dyn Module,
-    unwind_context: &mut UnwindContext,
     sig: Signature,
     wrapper_name: &str,
     callee_name: &str,
@@ -280,7 +279,6 @@ pub(crate) fn create_wrapper_function(
         bcx.finalize();
     }
     module.define_function(wrapper_func_id, &mut ctx).unwrap();
-    unwind_context.add_function(wrapper_func_id, &ctx, module.isa());
 }
 
 pub(crate) struct FunctionCx<'m, 'clif, 'tcx: 'm> {

src/driver/aot.rs

@@ -26,6 +26,7 @@
 use crate::concurrency_limiter::{ConcurrencyLimiter, ConcurrencyLimiterToken};
 use crate::debuginfo::TypeDebugContext;
 use crate::global_asm::GlobalAsmConfig;
+use crate::unwind_module::UnwindModule;
 use crate::{prelude::*, BackendConfig};
 
 struct ModuleCodegenResult {
@@ -318,7 +319,11 @@ fn produce_final_output_artifacts(
     // These are used in linking steps and will be cleaned up afterward.
 }
 
-fn make_module(sess: &Session, backend_config: &BackendConfig, name: String) -> ObjectModule {
+fn make_module(
+    sess: &Session,
+    backend_config: &BackendConfig,
+    name: String,
+) -> UnwindModule<ObjectModule> {
     let isa = crate::build_isa(sess, backend_config);
 
     let mut builder =
@@ -327,16 +332,15 @@ fn make_module(sess: &Session, backend_config: &BackendConfig, name: String) ->
     // is important, while cg_clif cares more about compilation times. Enabling -Zfunction-sections
     // can easily double the amount of time necessary to perform linking.
     builder.per_function_section(sess.opts.unstable_opts.function_sections.unwrap_or(false));
-    ObjectModule::new(builder)
+    UnwindModule::new(ObjectModule::new(builder), true)
 }
 
 fn emit_cgu(
     output_filenames: &OutputFilenames,
     prof: &SelfProfilerRef,
     name: String,
-    module: ObjectModule,
+    module: UnwindModule<ObjectModule>,
     debug: Option<DebugContext>,
-    unwind_context: UnwindContext,
     global_asm_object_file: Option<PathBuf>,
     producer: &str,
 ) -> Result<ModuleCodegenResult, String> {
@@ -346,8 +350,6 @@ fn emit_cgu(
         debug.emit(&mut product);
     }
 
-    unwind_context.emit(&mut product);
-
     let module_regular = emit_module(
         output_filenames,
         prof,
@@ -494,7 +496,6 @@ fn module_codegen(
 
             let mut cx = crate::CodegenCx::new(
                 tcx,
-                backend_config.clone(),
                 module.isa(),
                 tcx.sess.opts.debuginfo != DebugInfo::None,
                 cgu_name,
@@ -531,13 +532,7 @@ fn module_codegen(
                     }
                 }
             }
-            crate::main_shim::maybe_create_entry_wrapper(
+            crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, false, cgu.is_primary());
-                tcx,
-                &mut module,
-                &mut cx.unwind_context,
-                false,
-                cgu.is_primary(),
-            );
 
             let cgu_name = cgu.name().as_str().to_owned();
 
@@ -571,7 +566,6 @@ fn module_codegen(
                     cgu_name,
                     module,
                     cx.debug_context,
-                    cx.unwind_context,
                     global_asm_object_file,
                     &producer,
                 )
@@ -665,13 +659,10 @@ pub(crate) fn run_aot(
     });
 
     let mut allocator_module = make_module(tcx.sess, &backend_config, "allocator_shim".to_string());
-    let mut allocator_unwind_context = UnwindContext::new(allocator_module.isa(), true);
+    let created_alloc_shim = crate::allocator::codegen(tcx, &mut allocator_module);
-    let created_alloc_shim =
-        crate::allocator::codegen(tcx, &mut allocator_module, &mut allocator_unwind_context);
 
     let allocator_module = if created_alloc_shim {
-        let mut product = allocator_module.finish();
+        let product = allocator_module.finish();
-        allocator_unwind_context.emit(&mut product);
 
         match emit_module(
             tcx.output_filenames(()),

src/driver/jit.rs

@@ -14,12 +14,12 @@
 use rustc_span::Symbol;
 
 use crate::debuginfo::TypeDebugContext;
+use crate::unwind_module::UnwindModule;
 use crate::{prelude::*, BackendConfig};
 use crate::{CodegenCx, CodegenMode};
 
 struct JitState {
-    backend_config: BackendConfig,
+    jit_module: UnwindModule<JITModule>,
-    jit_module: JITModule,
 }
 
 thread_local! {
@@ -63,7 +63,7 @@ fn create_jit_module(
     tcx: TyCtxt<'_>,
     backend_config: &BackendConfig,
     hotswap: bool,
-) -> (JITModule, CodegenCx) {
+) -> (UnwindModule<JITModule>, CodegenCx) {
     let crate_info = CrateInfo::new(tcx, "dummy_target_cpu".to_string());
 
     let isa = crate::build_isa(tcx.sess, backend_config);
@@ -72,17 +72,11 @@ fn create_jit_module(
     crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
     jit_builder.symbol_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info));
     jit_builder.symbol("__clif_jit_fn", clif_jit_fn as *const u8);
-    let mut jit_module = JITModule::new(jit_builder);
+    let mut jit_module = UnwindModule::new(JITModule::new(jit_builder), false);
 
-    let mut cx = crate::CodegenCx::new(
+    let cx = crate::CodegenCx::new(tcx, jit_module.isa(), false, Symbol::intern("dummy_cgu_name"));
-        tcx,
-        backend_config.clone(),
-        jit_module.isa(),
-        false,
-        Symbol::intern("dummy_cgu_name"),
-    );
 
-    crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
+    crate::allocator::codegen(tcx, &mut jit_module);
 
     (jit_module, cx)
 }
@@ -128,7 +122,7 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
                         );
                     }
                     CodegenMode::JitLazy => {
-                        codegen_shim(tcx, &mut cx, &mut cached_context, &mut jit_module, inst)
+                        codegen_shim(tcx, &mut cached_context, &mut jit_module, inst)
                     }
                 },
                 MonoItem::Static(def_id) => {
@@ -146,18 +140,11 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         tcx.dcx().fatal("Inline asm is not supported in JIT mode");
     }
 
-    crate::main_shim::maybe_create_entry_wrapper(
+    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, true, true);
-        tcx,
-        &mut jit_module,
-        &mut cx.unwind_context,
-        true,
-        true,
-    );
 
     tcx.dcx().abort_if_errors();
 
-    jit_module.finalize_definitions().unwrap();
+    jit_module.finalize_definitions();
-    unsafe { cx.unwind_context.register_jit(&jit_module) };
 
     println!(
         "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
@@ -177,12 +164,12 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         call_conv: jit_module.target_config().default_call_conv,
     };
     let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
-    let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);
+    let finalized_start: *const u8 = jit_module.module.get_finalized_function(start_func_id);
 
     LAZY_JIT_STATE.with(|lazy_jit_state| {
         let mut lazy_jit_state = lazy_jit_state.borrow_mut();
         assert!(lazy_jit_state.is_none());
-        *lazy_jit_state = Some(JitState { backend_config, jit_module });
+        *lazy_jit_state = Some(JitState { jit_module });
     });
 
     let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
@@ -268,7 +255,6 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
             let mut lazy_jit_state = lazy_jit_state.borrow_mut();
             let lazy_jit_state = lazy_jit_state.as_mut().unwrap();
             let jit_module = &mut lazy_jit_state.jit_module;
-            let backend_config = lazy_jit_state.backend_config.clone();
 
             let name = tcx.symbol_name(instance).name;
             let sig = crate::abi::get_function_sig(
@@ -278,7 +264,7 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
             );
             let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
 
-            let current_ptr = jit_module.read_got_entry(func_id);
+            let current_ptr = jit_module.module.read_got_entry(func_id);
 
             // If the function's GOT entry has already been updated to point at something other
             // than the shim trampoline, don't re-jit but just return the new pointer instead.
@@ -288,11 +274,10 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
                 return current_ptr;
             }
 
-            jit_module.prepare_for_function_redefine(func_id).unwrap();
+            jit_module.module.prepare_for_function_redefine(func_id).unwrap();
 
             let mut cx = crate::CodegenCx::new(
                 tcx,
-                backend_config,
                 jit_module.isa(),
                 false,
                 Symbol::intern("dummy_cgu_name"),
@@ -300,9 +285,8 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
             codegen_and_compile_fn(tcx, &mut cx, &mut Context::new(), jit_module, instance);
 
             assert!(cx.global_asm.is_empty());
-            jit_module.finalize_definitions().unwrap();
+            jit_module.finalize_definitions();
-            unsafe { cx.unwind_context.register_jit(&jit_module) };
+            jit_module.module.get_finalized_function(func_id)
-            jit_module.get_finalized_function(func_id)
         })
     })
 }
@@ -362,9 +346,8 @@ fn dep_symbol_lookup_fn(
 
 fn codegen_shim<'tcx>(
     tcx: TyCtxt<'tcx>,
-    cx: &mut CodegenCx,
     cached_context: &mut Context,
-    module: &mut JITModule,
+    module: &mut UnwindModule<JITModule>,
     inst: Instance<'tcx>,
 ) {
     let pointer_type = module.target_config().pointer_type();
@@ -413,5 +396,4 @@ fn codegen_shim<'tcx>(
     trampoline_builder.ins().return_(&ret_vals);
 
     module.define_function(func_id, context).unwrap();
-    cx.unwind_context.add_function(func_id, context, module.isa());
 }

src/inline_asm.rs

@@ -113,13 +113,7 @@ pub(crate) fn codegen_inline_asm_terminator<'tcx>(
                     );
                     let sig =
                         get_function_sig(fx.tcx, fx.target_config.default_call_conv, instance);
-                    create_wrapper_function(
+                    create_wrapper_function(fx.module, sig, &wrapper_name, symbol.name);
-                        fx.module,
-                        &mut fx.cx.unwind_context,
-                        sig,
-                        &wrapper_name,
-                        symbol.name,
-                    );
 
                     CInlineAsmOperand::Symbol { symbol: wrapper_name }
                 } else {
@@ -283,13 +277,7 @@ pub(crate) fn codegen_naked_asm<'tcx>(
                     );
                     let sig =
                         get_function_sig(tcx, module.target_config().default_call_conv, instance);
-                    create_wrapper_function(
+                    create_wrapper_function(module, sig, &wrapper_name, symbol.name);
-                        module,
-                        &mut cx.unwind_context,
-                        sig,
-                        &wrapper_name,
-                        symbol.name,
-                    );
 
                     CInlineAsmOperand::Symbol { symbol: wrapper_name }
                 } else {

src/lib.rs

@@ -79,6 +79,7 @@
 mod toolchain;
 mod trap;
 mod unsize;
+mod unwind_module;
 mod value_and_place;
 mod vtable;
 
@@ -130,22 +131,13 @@ struct CodegenCx {
     global_asm: String,
     inline_asm_index: Cell<usize>,
     debug_context: Option<DebugContext>,
-    unwind_context: UnwindContext,
     cgu_name: Symbol,
 }
 
 impl CodegenCx {
-    fn new(
+    fn new(tcx: TyCtxt<'_>, isa: &dyn TargetIsa, debug_info: bool, cgu_name: Symbol) -> Self {
-        tcx: TyCtxt<'_>,
-        backend_config: BackendConfig,
-        isa: &dyn TargetIsa,
-        debug_info: bool,
-        cgu_name: Symbol,
-    ) -> Self {
         assert_eq!(pointer_ty(tcx), isa.pointer_type());
 
-        let unwind_context =
-            UnwindContext::new(isa, matches!(backend_config.codegen_mode, CodegenMode::Aot));
         let debug_context = if debug_info && !tcx.sess.target.options.is_like_windows {
             Some(DebugContext::new(tcx, isa, cgu_name.as_str()))
         } else {
@@ -158,7 +150,6 @@ fn new(
             global_asm: String::new(),
             inline_asm_index: Cell::new(0),
             debug_context,
-            unwind_context,
             cgu_name,
         }
     }

src/main_shim.rs

@@ -12,7 +12,6 @@
 pub(crate) fn maybe_create_entry_wrapper(
     tcx: TyCtxt<'_>,
     module: &mut impl Module,
-    unwind_context: &mut UnwindContext,
     is_jit: bool,
     is_primary_cgu: bool,
 ) {
@@ -36,12 +35,11 @@ pub(crate) fn maybe_create_entry_wrapper(
         return;
     }
 
-    create_entry_fn(tcx, module, unwind_context, main_def_id, is_jit, is_main_fn, sigpipe);
+    create_entry_fn(tcx, module, main_def_id, is_jit, is_main_fn, sigpipe);
 
     fn create_entry_fn(
         tcx: TyCtxt<'_>,
         m: &mut impl Module,
-        unwind_context: &mut UnwindContext,
         rust_main_def_id: DefId,
         ignore_lang_start_wrapper: bool,
         is_main_fn: bool,
@@ -170,7 +168,5 @@ fn create_entry_fn(
         if let Err(err) = m.define_function(cmain_func_id, &mut ctx) {
             tcx.dcx().fatal(format!("entry symbol `{entry_name}` defined multiple times: {err}"));
         }
-
-        unwind_context.add_function(cmain_func_id, &ctx, m.isa());
     }
 }

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)
+    }
+}