Move LTO symbol export calculation from backends to cg_ssa

3 files changed, 110 insertions, 0 deletions

diff --git a/compiler/rustc_codegen_ssa/messages.ftl b/compiler/rustc_codegen_ssa/messages.ftl
index c7bd6ffd1f2..a70d0011d16 100644
--- a/compiler/rustc_codegen_ssa/messages.ftl
+++ b/compiler/rustc_codegen_ssa/messages.ftl
@@ -35,6 +35,10 @@ codegen_ssa_dlltool_fail_import_library =
     {$stdout}
     {$stderr}
 
+codegen_ssa_dynamic_linking_with_lto =
+    cannot prefer dynamic linking when performing LTO
+    .note = only 'staticlib', 'bin', and 'cdylib' outputs are supported with LTO
+
 codegen_ssa_error_calling_dlltool =
     Error calling dlltool '{$dlltool_path}': {$error}
 
@@ -191,6 +195,12 @@ codegen_ssa_linker_unsupported_modifier = `as-needed` modifier not supported for
 
 codegen_ssa_linking_failed = linking with `{$linker_path}` failed: {$exit_status}
 
+codegen_ssa_lto_disallowed = lto can only be run for executables, cdylibs and static library outputs
+
+codegen_ssa_lto_dylib = lto cannot be used for `dylib` crate type without `-Zdylib-lto`
+
+codegen_ssa_lto_proc_macro = lto cannot be used for `proc-macro` crate type without `-Zdylib-lto`
+
 codegen_ssa_malformed_cgu_name =
     found malformed codegen unit name `{$user_path}`. codegen units names must always start with the name of the crate (`{$crate_name}` in this case).
 
diff --git a/compiler/rustc_codegen_ssa/src/back/lto.rs b/compiler/rustc_codegen_ssa/src/back/lto.rs
index b49b6783bbd..474a5e08baf 100644
--- a/compiler/rustc_codegen_ssa/src/back/lto.rs
+++ b/compiler/rustc_codegen_ssa/src/back/lto.rs
@@ -2,7 +2,15 @@ use std::ffi::CString;
 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_session::config::{CrateType, Lto};
+use tracing::info;
 
+use crate::back::symbol_export;
+use crate::back::write::CodegenContext;
+use crate::errors::{DynamicLinkingWithLTO, LtoDisallowed, LtoDylib, LtoProcMacro};
 use crate::traits::*;
 
 pub struct ThinModule<B: WriteBackendMethods> {
@@ -52,3 +60,78 @@ impl<M: ModuleBufferMethods> SerializedModule<M> {
         }
     }
 }
+
+fn crate_type_allows_lto(crate_type: CrateType) -> bool {
+    match crate_type {
+        CrateType::Executable
+        | CrateType::Dylib
+        | CrateType::Staticlib
+        | CrateType::Cdylib
+        | CrateType::ProcMacro
+        | CrateType::Sdylib => true,
+        CrateType::Rlib => false,
+    }
+}
+
+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 {
+        // 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::No => panic!("didn't request LTO but we're doing LTO"),
+    };
+
+    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 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>>()
+    };
+    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 cgcx.opts.cg.prefer_dynamic && !cgcx.opts.unstable_opts.dylib_lto {
+            return Err(dcx.emit_almost_fatal(DynamicLinkingWithLTO));
+        }
+
+        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));
+            }
+        }
+    }
+
+    Ok(symbols_below_threshold)
+}
diff --git a/compiler/rustc_codegen_ssa/src/errors.rs b/compiler/rustc_codegen_ssa/src/errors.rs
index 9040915b6af..3d787d8bdbd 100644
--- a/compiler/rustc_codegen_ssa/src/errors.rs
+++ b/compiler/rustc_codegen_ssa/src/errors.rs
@@ -1294,3 +1294,20 @@ pub(crate) struct FeatureNotValid<'a> {
     #[help]
     pub plus_hint: bool,
 }
+
+#[derive(Diagnostic)]
+#[diag(codegen_ssa_lto_disallowed)]
+pub(crate) struct LtoDisallowed;
+
+#[derive(Diagnostic)]
+#[diag(codegen_ssa_lto_dylib)]
+pub(crate) struct LtoDylib;
+
+#[derive(Diagnostic)]
+#[diag(codegen_ssa_lto_proc_macro)]
+pub(crate) struct LtoProcMacro;
+
+#[derive(Diagnostic)]
+#[diag(codegen_ssa_dynamic_linking_with_lto)]
+#[note]
+pub(crate) struct DynamicLinkingWithLTO;