Merge exported_symbols computation into exported_symbols_for_lto

bjorn3 · bjorn3 · commit 097a84f6ceae · 2025-07-17T16:45:09.000Z
And move exported_symbols_for_lto call from backends to cg_ssa.
diff --git a/compiler/rustc_codegen_gcc/src/back/lto.rs b/compiler/rustc_codegen_gcc/src/back/lto.rs
@@ -24,9 +24,7 @@ use std::sync::Arc;
 
 use gccjit::{Context, OutputKind};
 use object::read::archive::ArchiveFile;
-use rustc_codegen_ssa::back::lto::{
-    SerializedModule, ThinModule, ThinShared, exported_symbols_for_lto,
-};
+use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
 use rustc_codegen_ssa::traits::*;
 use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
@@ -54,7 +52,7 @@ fn prepare_lto(
     dcx: DiagCtxtHandle<'_>,
 ) -> Result<LtoData, FatalError> {
     // FIXME(bjorn3): Limit LTO exports to these symbols
-    let _symbols_below_threshold = exported_symbols_for_lto(cgcx, dcx)?;
+    let _symbols_below_threshold = &cgcx.exported_symbols_for_lto;
 
     let tmp_path = match tempdir() {
         Ok(tmp_path) => tmp_path,
diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -7,9 +7,7 @@ use std::sync::Arc;
 use std::{io, iter, slice};
 
 use object::read::archive::ArchiveFile;
-use rustc_codegen_ssa::back::lto::{
-    SerializedModule, ThinModule, ThinShared, exported_symbols_for_lto,
-};
+use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
 use rustc_codegen_ssa::traits::*;
 use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
@@ -37,9 +35,10 @@ fn prepare_lto(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     dcx: DiagCtxtHandle<'_>,
 ) -> Result<(Vec<CString>, Vec<(SerializedModule<ModuleBuffer>, CString)>), FatalError> {
-    let mut symbols_below_threshold = exported_symbols_for_lto(cgcx, dcx)?
-        .into_iter()
-        .map(|symbol| CString::new(symbol).unwrap())
+    let mut symbols_below_threshold = cgcx
+        .exported_symbols_for_lto
+        .iter()
+        .map(|symbol| CString::new(symbol.to_owned()).unwrap())
         .collect::<Vec<CString>>();
 
     // __llvm_profile_counter_bias is pulled in at link time by an undefined reference to
diff --git a/compiler/rustc_codegen_ssa/src/back/lto.rs b/compiler/rustc_codegen_ssa/src/back/lto.rs
@@ -1,14 +1,15 @@
 use std::ffi::CString;
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use rustc_data_structures::memmap::Mmap;
-use rustc_errors::{DiagCtxtHandle, FatalError};
-use rustc_hir::def_id::LOCAL_CRATE;
-use rustc_middle::middle::exported_symbols::{SymbolExportInfo, SymbolExportLevel};
+use rustc_hir::def_id::{CrateNum, LOCAL_CRATE};
+use rustc_middle::middle::exported_symbols::{ExportedSymbol, SymbolExportInfo, SymbolExportLevel};
+use rustc_middle::ty::TyCtxt;
 use rustc_session::config::{CrateType, Lto};
 use tracing::info;
 
-use crate::back::symbol_export;
+use crate::back::symbol_export::{self, symbol_name_for_instance_in_crate};
 use crate::back::write::CodegenContext;
 use crate::errors::{DynamicLinkingWithLTO, LtoDisallowed, LtoDylib, LtoProcMacro};
 use crate::traits::*;
@@ -73,65 +74,73 @@ fn crate_type_allows_lto(crate_type: CrateType) -> bool {
     }
 }
 
-pub fn exported_symbols_for_lto<'a, B: WriteBackendMethods>(
-    cgcx: &'a CodegenContext<B>,
-    dcx: DiagCtxtHandle<'_>,
-) -> Result<Vec<&'a str>, FatalError> {
-    // FIXME move symbol filtering to cg_ssa
-    let export_threshold = match cgcx.lto {
+pub(super) fn exported_symbols_for_lto(
+    tcx: TyCtxt<'_>,
+    each_linked_rlib_for_lto: &[(CrateNum, PathBuf)],
+) -> Vec<String> {
+    let export_threshold = match tcx.sess.lto() {
         // We're just doing LTO for our one crate
         Lto::ThinLocal => SymbolExportLevel::Rust,
 
         // We're doing LTO for the entire crate graph
-        Lto::Fat | Lto::Thin => symbol_export::crates_export_threshold(&cgcx.crate_types),
+        Lto::Fat | Lto::Thin => symbol_export::crates_export_threshold(&tcx.crate_types()),
 
-        Lto::No => panic!("didn't request LTO but we're doing LTO"),
+        Lto::No => return vec![],
     };
 
-    let symbol_filter = &|&(ref name, info): &'a (String, SymbolExportInfo)| {
-        if info.level.is_below_threshold(export_threshold) || info.used {
-            Some(name.as_str())
-        } else {
-            None
-        }
+    let copy_symbols = |cnum| {
+        tcx.exported_non_generic_symbols(cnum)
+            .iter()
+            .chain(tcx.exported_generic_symbols(cnum))
+            .filter_map(|&(s, info): &(ExportedSymbol<'_>, SymbolExportInfo)| {
+                if info.level.is_below_threshold(export_threshold) || info.used {
+                    Some(symbol_name_for_instance_in_crate(tcx, s, cnum))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>()
     };
-    let exported_symbols = cgcx.exported_symbols.as_ref().expect("needs exported symbols for LTO");
     let mut symbols_below_threshold = {
-        let _timer = cgcx.prof.generic_activity("lto_generate_symbols_below_threshold");
-        exported_symbols[&LOCAL_CRATE].iter().filter_map(symbol_filter).collect::<Vec<&str>>()
+        let _timer = tcx.prof.generic_activity("lto_generate_symbols_below_threshold");
+        copy_symbols(LOCAL_CRATE)
     };
     info!("{} symbols to preserve in this crate", symbols_below_threshold.len());
 
     // If we're performing LTO for the entire crate graph, then for each of our
     // upstream dependencies, include their exported symbols.
-    if cgcx.lto != Lto::ThinLocal {
-        // Make sure we actually can run LTO
-        for crate_type in cgcx.crate_types.iter() {
-            if !crate_type_allows_lto(*crate_type) {
-                return Err(dcx.emit_almost_fatal(LtoDisallowed));
-            } else if *crate_type == CrateType::Dylib {
-                if !cgcx.opts.unstable_opts.dylib_lto {
-                    return Err(dcx.emit_almost_fatal(LtoDylib));
-                }
-            } else if *crate_type == CrateType::ProcMacro && !cgcx.opts.unstable_opts.dylib_lto {
-                return Err(dcx.emit_almost_fatal(LtoProcMacro));
-            }
+    if tcx.sess.lto() != Lto::ThinLocal {
+        for &(cnum, ref _path) in each_linked_rlib_for_lto {
+            let _timer = tcx.prof.generic_activity("lto_generate_symbols_below_threshold");
+            symbols_below_threshold.extend(copy_symbols(cnum));
         }
+    }
 
-        if cgcx.opts.cg.prefer_dynamic && !cgcx.opts.unstable_opts.dylib_lto {
-            return Err(dcx.emit_almost_fatal(DynamicLinkingWithLTO));
-        }
+    symbols_below_threshold
+}
+
+pub(super) fn check_lto_allowed<B: WriteBackendMethods>(cgcx: &CodegenContext<B>) {
+    if cgcx.lto == Lto::ThinLocal {
+        // Crate local LTO is always allowed
+        return;
+    }
 
-        for &(cnum, ref _path) in cgcx.each_linked_rlib_for_lto.iter() {
-            let exported_symbols =
-                cgcx.exported_symbols.as_ref().expect("needs exported symbols for LTO");
-            {
-                let _timer = cgcx.prof.generic_activity("lto_generate_symbols_below_threshold");
-                symbols_below_threshold
-                    .extend(exported_symbols[&cnum].iter().filter_map(symbol_filter));
+    let dcx = cgcx.create_dcx();
+
+    // Make sure we actually can run LTO
+    for crate_type in cgcx.crate_types.iter() {
+        if !crate_type_allows_lto(*crate_type) {
+            dcx.handle().emit_fatal(LtoDisallowed);
+        } else if *crate_type == CrateType::Dylib {
+            if !cgcx.opts.unstable_opts.dylib_lto {
+                dcx.handle().emit_fatal(LtoDylib);
             }
+        } else if *crate_type == CrateType::ProcMacro && !cgcx.opts.unstable_opts.dylib_lto {
+            dcx.handle().emit_fatal(LtoProcMacro);
         }
     }
 
-    Ok(symbols_below_threshold)
+    if cgcx.opts.cg.prefer_dynamic && !cgcx.opts.unstable_opts.dylib_lto {
+        dcx.handle().emit_fatal(DynamicLinkingWithLTO);
+    }
 }
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -9,7 +9,7 @@ use std::{fs, io, mem, str, thread};
 use rustc_abi::Size;
 use rustc_ast::attr;
 use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
-use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::jobserver::{self, Acquired};
 use rustc_data_structures::memmap::Mmap;
 use rustc_data_structures::profiling::{SelfProfilerRef, VerboseTimingGuard};
@@ -20,14 +20,13 @@ use rustc_errors::{
     Suggestions,
 };
 use rustc_fs_util::link_or_copy;
-use rustc_hir::def_id::{CrateNum, LOCAL_CRATE};
+use rustc_hir::def_id::CrateNum;
 use rustc_incremental::{
     copy_cgu_workproduct_to_incr_comp_cache_dir, in_incr_comp_dir, in_incr_comp_dir_sess,
 };
 use rustc_metadata::fs::copy_to_stdout;
 use rustc_middle::bug;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
-use rustc_middle::middle::exported_symbols::SymbolExportInfo;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::Session;
 use rustc_session::config::{
@@ -40,7 +39,7 @@ use tracing::debug;
 
 use super::link::{self, ensure_removed};
 use super::lto::{self, SerializedModule};
-use super::symbol_export::symbol_name_for_instance_in_crate;
+use crate::back::lto::check_lto_allowed;
 use crate::errors::{AutodiffWithoutLto, ErrorCreatingRemarkDir};
 use crate::traits::*;
 use crate::{
@@ -330,8 +329,6 @@ pub type TargetMachineFactoryFn<B> = Arc<
         + Sync,
 >;
 
-type ExportedSymbols = FxHashMap<CrateNum, Arc<Vec<(String, SymbolExportInfo)>>>;
-
 /// Additional resources used by optimize_and_codegen (not module specific)
 #[derive(Clone)]
 pub struct CodegenContext<B: WriteBackendMethods> {
@@ -341,10 +338,10 @@ pub struct CodegenContext<B: WriteBackendMethods> {
     pub save_temps: bool,
     pub fewer_names: bool,
     pub time_trace: bool,
-    pub exported_symbols: Option<Arc<ExportedSymbols>>,
     pub opts: Arc<config::Options>,
     pub crate_types: Vec<CrateType>,
     pub each_linked_rlib_for_lto: Vec<(CrateNum, PathBuf)>,
+    pub exported_symbols_for_lto: Arc<Vec<String>>,
     pub output_filenames: Arc<OutputFilenames>,
     pub invocation_temp: Option<String>,
     pub regular_module_config: Arc<ModuleConfig>,
@@ -1126,34 +1123,8 @@ fn start_executing_work<B: ExtraBackendMethods>(
     }));
 
     // Compute the set of symbols we need to retain when doing LTO (if we need to)
-    let exported_symbols = {
-        let mut exported_symbols = FxHashMap::default();
-
-        let copy_symbols = |cnum| {
-            let symbols = tcx
-                .exported_non_generic_symbols(cnum)
-                .iter()
-                .chain(tcx.exported_generic_symbols(cnum))
-                .map(|&(s, lvl)| (symbol_name_for_instance_in_crate(tcx, s, cnum), lvl))
-                .collect();
-            Arc::new(symbols)
-        };
-
-        match sess.lto() {
-            Lto::No => None,
-            Lto::ThinLocal => {
-                exported_symbols.insert(LOCAL_CRATE, copy_symbols(LOCAL_CRATE));
-                Some(Arc::new(exported_symbols))
-            }
-            Lto::Fat | Lto::Thin => {
-                exported_symbols.insert(LOCAL_CRATE, copy_symbols(LOCAL_CRATE));
-                for &(cnum, ref _path) in &each_linked_rlib_for_lto {
-                    exported_symbols.insert(cnum, copy_symbols(cnum));
-                }
-                Some(Arc::new(exported_symbols))
-            }
-        }
-    };
+    let exported_symbols_for_lto =
+        Arc::new(lto::exported_symbols_for_lto(tcx, &each_linked_rlib_for_lto));
 
     // First up, convert our jobserver into a helper thread so we can use normal
     // mpsc channels to manage our messages and such.
@@ -1189,13 +1160,13 @@ fn start_executing_work<B: ExtraBackendMethods>(
     let cgcx = CodegenContext::<B> {
         crate_types: tcx.crate_types().to_vec(),
         each_linked_rlib_for_lto,
+        exported_symbols_for_lto,
         lto: sess.lto(),
         fewer_names: sess.fewer_names(),
         save_temps: sess.opts.cg.save_temps,
         time_trace: sess.opts.unstable_opts.llvm_time_trace,
         opts: Arc::new(sess.opts.clone()),
         prof: sess.prof.clone(),
-        exported_symbols,
         remark: sess.opts.cg.remark.clone(),
         remark_dir,
         incr_comp_session_dir: sess.incr_comp_session_dir_opt().map(|r| r.clone()),
@@ -1463,6 +1434,8 @@ fn start_executing_work<B: ExtraBackendMethods>(
                     let needs_thin_lto = mem::take(&mut needs_thin_lto);
                     let import_only_modules = mem::take(&mut lto_import_only_modules);
 
+                    check_lto_allowed(&cgcx);
+
                     if !needs_fat_lto.is_empty() {
                         assert!(needs_thin_lto.is_empty());
 
