Auto merge of #55871 - ljedrz:llvm_back_allocations, r=nagisa

codegen_llvm_back: improve allocations

This commit was split out from https://github.com/rust-lang/rust/pull/54864. Last time it was causing an LLVM OOM, which was most probably caused by not collecting the globals.

- preallocate vectors of known length
- `extend` instead of `append` where the argument is consumable
- turn 2 `push` loops into `extend`s
- create a vector from a function producing one instead of using `extend_from_slice` on it
- consume `modules` when no longer needed
- ~~return an `impl Iterator` from `generate_lto_work`~~
- ~~don't `collect` `globals`, as they are iterated over and consumed right afterwards~~

While I'm hoping it won't cause an OOM anymore, I would still consider this a "high-risk" PR and not roll it up.

4 files changed, 23 insertions, 21 deletions

diff --git a/src/librustc_codegen_llvm/back/link.rs b/src/librustc_codegen_llvm/back/link.rs
index 89d84acdd88..68da14570e4 100644
--- a/src/librustc_codegen_llvm/back/link.rs
+++ b/src/librustc_codegen_llvm/back/link.rs
@@ -72,12 +72,12 @@ pub(crate) fn link_binary(sess: &Session,
            bug!("invalid output type `{:?}` for target os `{}`",
                 crate_type, sess.opts.target_triple);
         }
-        let mut out_files = link_binary_output(sess,
-                                               codegen_results,
-                                               crate_type,
-                                               outputs,
-                                               crate_name);
-        out_filenames.append(&mut out_files);
+        let out_files = link_binary_output(sess,
+                                           codegen_results,
+                                           crate_type,
+                                           outputs,
+                                           crate_name);
+        out_filenames.extend(out_files);
     }
 
     // Remove the temporary object file and metadata if we aren't saving temps
diff --git a/src/librustc_codegen_llvm/back/lto.rs b/src/librustc_codegen_llvm/back/lto.rs
index b5ebd0409da..99828e5b7fb 100644
--- a/src/librustc_codegen_llvm/back/lto.rs
+++ b/src/librustc_codegen_llvm/back/lto.rs
@@ -225,11 +225,12 @@ fn fat_lto(cgcx: &CodegenContext<LlvmCodegenBackend>,
         // and we want to move everything to the same LLVM context. Currently the
         // way we know of to do that is to serialize them to a string and them parse
         // them later. Not great but hey, that's why it's "fat" LTO, right?
-        for module in modules {
+        serialized_modules.extend(modules.into_iter().map(|module| {
             let buffer = ModuleBuffer::new(module.module_llvm.llmod());
             let llmod_id = CString::new(&module.name[..]).unwrap();
-            serialized_modules.push((SerializedModule::Local(buffer), llmod_id));
-        }
+
+            (SerializedModule::Local(buffer), llmod_id)
+        }));
 
         // For all serialized bitcode files we parse them and link them in as we did
         // above, this is all mostly handled in C++. Like above, though, we don't
@@ -349,9 +350,10 @@ fn thin_lto(cgcx: &CodegenContext<LlvmCodegenBackend>,
             .map(|&(_, ref wp)| (wp.cgu_name.clone(), wp.clone()))
             .collect();
 
-        let mut thin_buffers = Vec::new();
-        let mut module_names = Vec::new();
-        let mut thin_modules = Vec::new();
+        let full_scope_len = modules.len() + serialized_modules.len() + cached_modules.len();
+        let mut thin_buffers = Vec::with_capacity(modules.len());
+        let mut module_names = Vec::with_capacity(full_scope_len);
+        let mut thin_modules = Vec::with_capacity(full_scope_len);
 
         // FIXME: right now, like with fat LTO, we serialize all in-memory
         //        modules before working with them and ThinLTO. We really
@@ -360,7 +362,7 @@ fn thin_lto(cgcx: &CodegenContext<LlvmCodegenBackend>,
         //        into the global index. It turns out that this loop is by far
         //        the most expensive portion of this small bit of global
         //        analysis!
-        for (i, module) in modules.iter().enumerate() {
+        for (i, module) in modules.into_iter().enumerate() {
             info!("local module: {} - {}", i, module.name);
             let name = CString::new(module.name.clone()).unwrap();
             let buffer = ThinBuffer::new(module.module_llvm.llmod());
@@ -406,7 +408,7 @@ fn thin_lto(cgcx: &CodegenContext<LlvmCodegenBackend>,
         //        incremental ThinLTO first where we could actually avoid
         //        looking at upstream modules entirely sometimes (the contents,
         //        we must always unconditionally look at the index).
-        let mut serialized = Vec::new();
+        let mut serialized = Vec::with_capacity(serialized_modules.len() + cached_modules.len());
 
         let cached_modules = cached_modules.into_iter().map(|(sm, wp)| {
             (sm, CString::new(wp.cgu_name).unwrap())
diff --git a/src/librustc_codegen_llvm/back/rpath.rs b/src/librustc_codegen_llvm/back/rpath.rs
index ee4a9b30a1a..73a7366d0a3 100644
--- a/src/librustc_codegen_llvm/back/rpath.rs
+++ b/src/librustc_codegen_llvm/back/rpath.rs
@@ -31,14 +31,12 @@ pub fn get_rpath_flags(config: &mut RPathConfig) -> Vec<String> {
         return Vec::new();
     }
 
-    let mut flags = Vec::new();
-
     debug!("preparing the RPATH!");
 
     let libs = config.used_crates.clone();
     let libs = libs.iter().filter_map(|&(_, ref l)| l.option()).collect::<Vec<_>>();
     let rpaths = get_rpaths(config, &libs);
-    flags.extend_from_slice(&rpaths_to_flags(&rpaths));
+    let mut flags = rpaths_to_flags(&rpaths);
 
     // Use DT_RUNPATH instead of DT_RPATH if available
     if config.linker_is_gnu {
@@ -49,7 +47,8 @@ pub fn get_rpath_flags(config: &mut RPathConfig) -> Vec<String> {
 }
 
 fn rpaths_to_flags(rpaths: &[String]) -> Vec<String> {
-    let mut ret = Vec::new();
+    let mut ret = Vec::with_capacity(rpaths.len()); // the minimum needed capacity
+
     for rpath in rpaths {
         if rpath.contains(',') {
             ret.push("-Wl,-rpath".into());
diff --git a/src/librustc_codegen_ssa/back/symbol_export.rs b/src/librustc_codegen_ssa/back/symbol_export.rs
index 0fb2641a4f8..d25917e51bf 100644
--- a/src/librustc_codegen_ssa/back/symbol_export.rs
+++ b/src/librustc_codegen_ssa/back/symbol_export.rs
@@ -229,10 +229,11 @@ fn exported_symbols_provider_local<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             "__llvm_profile_raw_version",
             "__llvm_profile_filename",
         ];
-        for sym in &PROFILER_WEAK_SYMBOLS {
+
+        symbols.extend(PROFILER_WEAK_SYMBOLS.iter().map(|sym| {
             let exported_symbol = ExportedSymbol::NoDefId(SymbolName::new(sym));
-            symbols.push((exported_symbol, SymbolExportLevel::C));
-        }
+            (exported_symbol, SymbolExportLevel::C)
+        }));
     }
 
     if tcx.sess.crate_types.borrow().contains(&config::CrateType::Dylib) {