Create JITModule and ObjectModule wrappers and move UnwindContext calls into them

Author: bjorn3

src/allocator.rs

@@ -11,15 +11,10 @@ use rustc_session::config::OomStrategy;
 use crate::prelude::*;
 
 /// Returns whether an allocator shim was created
-pub(crate) fn codegen(
-    tcx: TyCtxt<'_>,
-    module: &mut impl Module,
-    unwind_context: &mut UnwindContext,
-) -> bool {
+pub(crate) fn codegen(tcx: TyCtxt<'_>, module: &mut impl Module) -> 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

@@ -6,7 +6,12 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::thread::JoinHandle;
 
-use cranelift_object::{ObjectBuilder, ObjectModule};
+use cranelift_codegen::control::ControlPlane;
+use cranelift_codegen::ir::Signature;
+use cranelift_codegen::isa::{TargetFrontendConfig, TargetIsa};
+use cranelift_codegen::FinalizedMachReloc;
+use cranelift_module::ModuleResult;
+use cranelift_object::{ObjectBuilder, ObjectModule, ObjectProduct};
 use rustc_codegen_ssa::assert_module_sources::CguReuse;
 use rustc_codegen_ssa::back::link::ensure_removed;
 use rustc_codegen_ssa::back::metadata::create_compressed_metadata_file;
@@ -318,7 +323,91 @@ 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 {
+struct AOTModule {
+    aot_module: ObjectModule,
+    unwind_context: UnwindContext,
+}
+
+impl AOTModule {
+    fn finish(self) -> ObjectProduct {
+        let mut product = self.aot_module.finish();
+        self.unwind_context.emit(&mut product);
+        product
+    }
+}
+
+impl Module for AOTModule {
+    fn isa(&self) -> &dyn TargetIsa {
+        self.aot_module.isa()
+    }
+
+    fn declarations(&self) -> &cranelift_module::ModuleDeclarations {
+        self.aot_module.declarations()
+    }
+
+    fn get_name(&self, name: &str) -> Option<cranelift_module::FuncOrDataId> {
+        self.aot_module.get_name(name)
+    }
+
+    fn target_config(&self) -> TargetFrontendConfig {
+        self.aot_module.target_config()
+    }
+
+    fn declare_function(
+        &mut self,
+        name: &str,
+        linkage: Linkage,
+        signature: &Signature,
+    ) -> ModuleResult<FuncId> {
+        self.aot_module.declare_function(name, linkage, signature)
+    }
+
+    fn declare_anonymous_function(&mut self, signature: &Signature) -> ModuleResult<FuncId> {
+        self.aot_module.declare_anonymous_function(signature)
+    }
+
+    fn declare_data(
+        &mut self,
+        name: &str,
+        linkage: Linkage,
+        writable: bool,
+        tls: bool,
+    ) -> ModuleResult<DataId> {
+        self.aot_module.declare_data(name, linkage, writable, tls)
+    }
+
+    fn declare_anonymous_data(&mut self, writable: bool, tls: bool) -> ModuleResult<DataId> {
+        self.aot_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.aot_module.define_function_with_control_plane(func, ctx, ctrl_plane)?;
+        self.unwind_context.add_function(func, ctx, self.aot_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.aot_module.define_data(data_id, data)
+    }
+}
+
+fn make_module(sess: &Session, backend_config: &BackendConfig, name: String) -> AOTModule {
     let isa = crate::build_isa(sess, backend_config);
 
     let mut builder =
@@ -327,16 +416,17 @@ 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)
+    let aot_module = ObjectModule::new(builder);
+    let unwind_context = UnwindContext::new(aot_module.isa(), true);
+    AOTModule { aot_module, unwind_context }
 }
 
 fn emit_cgu(
     output_filenames: &OutputFilenames,
     prof: &SelfProfilerRef,
     name: String,
-    module: ObjectModule,
+    module: AOTModule,
     debug: Option<DebugContext>,
-    unwind_context: UnwindContext,
     global_asm_object_file: Option<PathBuf>,
     producer: &str,
 ) -> Result<ModuleCodegenResult, String> {
@@ -346,8 +436,6 @@ fn emit_cgu(
         debug.emit(&mut product);
     }
 
-    unwind_context.emit(&mut product);
-
     let module_regular = emit_module(
         output_filenames,
         prof,
@@ -494,7 +582,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 +618,7 @@ fn module_codegen(
                     }
                 }
             }
-            crate::main_shim::maybe_create_entry_wrapper(
-                tcx,
-                &mut module,
-                &mut cx.unwind_context,
-                false,
-                cgu.is_primary(),
-            );
+            crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, false, cgu.is_primary());
 
             let cgu_name = cgu.name().as_str().to_owned();
 
@@ -571,7 +652,6 @@ fn module_codegen(
                     cgu_name,
                     module,
                     cx.debug_context,
-                    cx.unwind_context,
                     global_asm_object_file,
                     &producer,
                 )
@@ -665,13 +745,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, &mut allocator_unwind_context);
+    let created_alloc_shim = crate::allocator::codegen(tcx, &mut allocator_module);
 
     let allocator_module = if created_alloc_shim {
-        let mut product = allocator_module.finish();
-        allocator_unwind_context.emit(&mut product);
+        let product = allocator_module.finish();
 
         match emit_module(
             tcx.output_filenames(()),