1 files changed, 91 insertions, 0 deletions
diff --git a/compiler/rustc_codegen_ssa/src/back/lto.rs b/compiler/rustc_codegen_ssa/src/back/lto.rs
index b49b6783bbd..c95038375a1 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_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::{self, symbol_name_for_instance_in_crate};
+use crate::back::write::CodegenContext;
+use crate::errors::{DynamicLinkingWithLTO, LtoDisallowed, LtoDylib, LtoProcMacro};
 use crate::traits::*;
 
 pub struct ThinModule<B: WriteBackendMethods> {
@@ -52,3 +60,86 @@ 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(super) fn exported_symbols_for_lto(
+    tcx: TyCtxt<'_>,
+    each_linked_rlib_for_lto: &[CrateNum],
+) -> 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(&tcx.crate_types()),
+
+        Lto::No => return vec![],
+    };
+
+    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 mut symbols_below_threshold = {
+        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 tcx.sess.lto() != Lto::ThinLocal {
+        for &cnum in each_linked_rlib_for_lto {
+            let _timer = tcx.prof.generic_activity("lto_generate_symbols_below_threshold");
+            symbols_below_threshold.extend(copy_symbols(cnum));
+        }
+    }
+
+    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;
+    }
+
+    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);
+        }
+    }
+
+    if cgcx.opts.cg.prefer_dynamic && !cgcx.opts.unstable_opts.dylib_lto {
+        dcx.handle().emit_fatal(DynamicLinkingWithLTO);
+    }
+}