Save pre-link bitcode to ModuleCodegen

Author: dianqk

Diff for: compiler/rustc_codegen_gcc/src/lib.rs

@@ -393,7 +393,7 @@ impl WriteBackendMethods for GccCodegenBackend {
     unsafe fn optimize(
         _cgcx: &CodegenContext<Self>,
         _dcx: DiagCtxtHandle<'_>,
-        module: &ModuleCodegen<Self::Module>,
+        module: &mut ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<(), FatalError> {
         module.module_llvm.context.set_optimization_level(to_gcc_opt_level(config.opt_level));

Diff for: compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -776,6 +776,10 @@ pub(crate) unsafe fn optimize_thin_module(
     // that LLVM Context and Module.
     let module_llvm = ModuleLlvm::parse(cgcx, module_name, thin_module.data(), dcx)?;
     let mut module = ModuleCodegen::new_regular(thin_module.name(), module_llvm);
+    // Given that the newly created module lacks a thinlto buffer for embedding, we need to re-add it here.
+    if cgcx.config(ModuleKind::Regular).embed_bitcode() {
+        module.thin_lto_buffer = Some(thin_module.data().to_vec());
+    }
     {
         let target = &*module.module_llvm.tm;
         let llmod = module.module_llvm.llmod();

Diff for: compiler/rustc_codegen_llvm/src/back/write.rs

@@ -689,7 +689,7 @@ pub(crate) unsafe fn llvm_optimize(
 pub(crate) unsafe fn optimize(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     dcx: DiagCtxtHandle<'_>,
-    module: &ModuleCodegen<ModuleLlvm>,
+    module: &mut ModuleCodegen<ModuleLlvm>,
     config: &ModuleConfig,
 ) -> Result<(), FatalError> {
     let _timer = cgcx.prof.generic_activity_with_arg("LLVM_module_optimize", &*module.name);
@@ -745,10 +745,7 @@ pub(crate) unsafe fn optimize(
         }?;
         if let Some(thin_lto_buffer) = thin_lto_buffer {
             let thin_lto_buffer = unsafe { ThinBuffer::from_raw_ptr(thin_lto_buffer) };
-            let thin_bc_out = cgcx.output_filenames.temp_path(OutputType::ThinBitcode, module_name);
-            if let Err(err) = fs::write(&thin_bc_out, thin_lto_buffer.data()) {
-                dcx.emit_err(WriteBytecode { path: &thin_bc_out, err });
-            }
+            module.thin_lto_buffer = Some(thin_lto_buffer.data().to_vec());
             let bc_summary_out =
                 cgcx.output_filenames.temp_path(OutputType::ThinLinkBitcode, module_name);
             if config.emit_thin_lto_summary
@@ -848,20 +845,14 @@ pub(crate) unsafe fn codegen(
                 }
             }
 
-            if config.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full)
-                && module.kind == ModuleKind::Regular
-            {
+            if config.embed_bitcode() && module.kind == ModuleKind::Regular {
                 let _timer = cgcx
                     .prof
                     .generic_activity_with_arg("LLVM_module_codegen_embed_bitcode", &*module.name);
-                let thin_bc_out =
-                    cgcx.output_filenames.temp_path(OutputType::ThinBitcode, module_name);
-                assert!(thin_bc_out.exists(), "cannot find {:?} as embedded bitcode", thin_bc_out);
-                let data = fs::read(&thin_bc_out).unwrap();
-                debug!("removing embed bitcode file {:?}", thin_bc_out);
-                ensure_removed(dcx, &thin_bc_out);
+                let thin_bc =
+                    module.thin_lto_buffer.as_deref().expect("cannot find embedded bitcode");
                 unsafe {
-                    embed_bitcode(cgcx, llcx, llmod, &config.bc_cmdline, &data);
+                    embed_bitcode(cgcx, llcx, llmod, &config.bc_cmdline, &thin_bc);
                 }
             }
         }

Diff for: compiler/rustc_codegen_llvm/src/lib.rs

@@ -194,7 +194,7 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     unsafe fn optimize(
         cgcx: &CodegenContext<Self>,
         dcx: DiagCtxtHandle<'_>,
-        module: &ModuleCodegen<Self::Module>,
+        module: &mut ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<(), FatalError> {
         unsafe { back::write::optimize(cgcx, dcx, module, config) }

Diff for: compiler/rustc_codegen_ssa/src/back/write.rs

@@ -278,6 +278,10 @@ impl ModuleConfig {
             || self.emit_obj == EmitObj::Bitcode
             || self.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full)
     }
+
+    pub fn embed_bitcode(&self) -> bool {
+        self.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full)
+    }
 }
 
 /// Configuration passed to the function returned by the `target_machine_factory`.
@@ -880,14 +884,14 @@ pub(crate) fn compute_per_cgu_lto_type(
 
 fn execute_optimize_work_item<B: ExtraBackendMethods>(
     cgcx: &CodegenContext<B>,
-    module: ModuleCodegen<B::Module>,
+    mut module: ModuleCodegen<B::Module>,
     module_config: &ModuleConfig,
 ) -> Result<WorkItemResult<B>, FatalError> {
     let dcx = cgcx.create_dcx();
     let dcx = dcx.handle();
 
     unsafe {
-        B::optimize(cgcx, dcx, &module, module_config)?;
+        B::optimize(cgcx, dcx, &mut module, module_config)?;
     }
 
     // After we've done the initial round of optimizations we need to

Diff for: compiler/rustc_codegen_ssa/src/lib.rs

@@ -75,15 +75,27 @@ pub struct ModuleCodegen<M> {
     pub name: String,
     pub module_llvm: M,
     pub kind: ModuleKind,
+    /// Saving the ThinLTO buffer for embedding in the object file.
+    pub thin_lto_buffer: Option<Vec<u8>>,
 }
 
 impl<M> ModuleCodegen<M> {
     pub fn new_regular(name: impl Into<String>, module: M) -> Self {
-        Self { name: name.into(), module_llvm: module, kind: ModuleKind::Regular }
+        Self {
+            name: name.into(),
+            module_llvm: module,
+            kind: ModuleKind::Regular,
+            thin_lto_buffer: None,
+        }
     }
 
     pub fn new_allocator(name: impl Into<String>, module: M) -> Self {
-        Self { name: name.into(), module_llvm: module, kind: ModuleKind::Allocator }
+        Self {
+            name: name.into(),
+            module_llvm: module,
+            kind: ModuleKind::Allocator,
+            thin_lto_buffer: None,
+        }
     }
 
     pub fn into_compiled_module(

Diff for: compiler/rustc_codegen_ssa/src/traits/write.rs

@@ -40,7 +40,7 @@ pub trait WriteBackendMethods: 'static + Sized + Clone {
     unsafe fn optimize(
         cgcx: &CodegenContext<Self>,
         dcx: DiagCtxtHandle<'_>,
-        module: &ModuleCodegen<Self::Module>,
+        module: &mut ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<(), FatalError>;
     fn optimize_fat(