1 files changed, 295 insertions, 0 deletions
diff --git a/compiler/rustc_codegen_ssa/src/assert_module_sources.rs b/compiler/rustc_codegen_ssa/src/assert_module_sources.rs
new file mode 100644
index 00000000000..16bb7b12bd3
--- /dev/null
+++ b/compiler/rustc_codegen_ssa/src/assert_module_sources.rs
@@ -0,0 +1,295 @@
+//! This pass is only used for UNIT TESTS related to incremental
+//! compilation. It tests whether a particular `.o` file will be re-used
+//! from a previous compilation or whether it must be regenerated.
+//!
+//! The user adds annotations to the crate of the following form:
+//!
+//! ```
+//! # #![feature(rustc_attrs)]
+//! # #![allow(internal_features)]
+//! #![rustc_partition_reused(module="spike", cfg="rpass2")]
+//! #![rustc_partition_codegened(module="spike-x", cfg="rpass2")]
+//! ```
+//!
+//! The first indicates (in the cfg `rpass2`) that `spike.o` will be
+//! reused, the second that `spike-x.o` will be recreated. If these
+//! annotations are inaccurate, errors are reported.
+//!
+//! The reason that we use `cfg=...` and not `#[cfg_attr]` is so that
+//! the HIR doesn't change as a result of the annotations, which might
+//! perturb the reuse results.
+//!
+//! `#![rustc_expected_cgu_reuse(module="spike", cfg="rpass2", kind="post-lto")]`
+//! allows for doing a more fine-grained check to see if pre- or post-lto data
+//! was re-used.
+
+use crate::errors;
+use rustc_ast as ast;
+use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::unord::UnordSet;
+use rustc_errors::{DiagnosticArgValue, IntoDiagnosticArg};
+use rustc_hir::def_id::LOCAL_CRATE;
+use rustc_middle::mir::mono::CodegenUnitNameBuilder;
+use rustc_middle::ty::TyCtxt;
+use rustc_session::Session;
+use rustc_span::symbol::sym;
+use rustc_span::{Span, Symbol};
+use std::borrow::Cow;
+use std::fmt;
+use thin_vec::ThinVec;
+
+#[allow(missing_docs)]
+pub fn assert_module_sources(tcx: TyCtxt<'_>, set_reuse: &dyn Fn(&mut CguReuseTracker)) {
+    tcx.dep_graph.with_ignore(|| {
+        if tcx.sess.opts.incremental.is_none() {
+            return;
+        }
+
+        let available_cgus =
+            tcx.collect_and_partition_mono_items(()).1.iter().map(|cgu| cgu.name()).collect();
+
+        let mut ams = AssertModuleSource {
+            tcx,
+            available_cgus,
+            cgu_reuse_tracker: if tcx.sess.opts.unstable_opts.query_dep_graph {
+                CguReuseTracker::new()
+            } else {
+                CguReuseTracker::new_disabled()
+            },
+        };
+
+        for attr in tcx.hir().attrs(rustc_hir::CRATE_HIR_ID) {
+            ams.check_attr(attr);
+        }
+
+        set_reuse(&mut ams.cgu_reuse_tracker);
+
+        ams.cgu_reuse_tracker.check_expected_reuse(tcx.sess);
+    });
+}
+
+struct AssertModuleSource<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    available_cgus: UnordSet<Symbol>,
+    cgu_reuse_tracker: CguReuseTracker,
+}
+
+impl<'tcx> AssertModuleSource<'tcx> {
+    fn check_attr(&mut self, attr: &ast::Attribute) {
+        let (expected_reuse, comp_kind) = if attr.has_name(sym::rustc_partition_reused) {
+            (CguReuse::PreLto, ComparisonKind::AtLeast)
+        } else if attr.has_name(sym::rustc_partition_codegened) {
+            (CguReuse::No, ComparisonKind::Exact)
+        } else if attr.has_name(sym::rustc_expected_cgu_reuse) {
+            match self.field(attr, sym::kind) {
+                sym::no => (CguReuse::No, ComparisonKind::Exact),
+                sym::pre_dash_lto => (CguReuse::PreLto, ComparisonKind::Exact),
+                sym::post_dash_lto => (CguReuse::PostLto, ComparisonKind::Exact),
+                sym::any => (CguReuse::PreLto, ComparisonKind::AtLeast),
+                other => {
+                    self.tcx
+                        .sess
+                        .emit_fatal(errors::UnknownReuseKind { span: attr.span, kind: other });
+                }
+            }
+        } else {
+            return;
+        };
+
+        if !self.tcx.sess.opts.unstable_opts.query_dep_graph {
+            self.tcx.sess.emit_fatal(errors::MissingQueryDepGraph { span: attr.span });
+        }
+
+        if !self.check_config(attr) {
+            debug!("check_attr: config does not match, ignoring attr");
+            return;
+        }
+
+        let user_path = self.field(attr, sym::module).to_string();
+        let crate_name = self.tcx.crate_name(LOCAL_CRATE).to_string();
+
+        if !user_path.starts_with(&crate_name) {
+            self.tcx.sess.emit_fatal(errors::MalformedCguName {
+                span: attr.span,
+                user_path,
+                crate_name,
+            });
+        }
+
+        // Split of the "special suffix" if there is one.
+        let (user_path, cgu_special_suffix) = if let Some(index) = user_path.rfind('.') {
+            (&user_path[..index], Some(&user_path[index + 1..]))
+        } else {
+            (&user_path[..], None)
+        };
+
+        let mut iter = user_path.split('-');
+
+        // Remove the crate name
+        assert_eq!(iter.next().unwrap(), crate_name);
+
+        let cgu_path_components = iter.collect::<Vec<_>>();
+
+        let cgu_name_builder = &mut CodegenUnitNameBuilder::new(self.tcx);
+        let cgu_name =
+            cgu_name_builder.build_cgu_name(LOCAL_CRATE, cgu_path_components, cgu_special_suffix);
+
+        debug!("mapping '{}' to cgu name '{}'", self.field(attr, sym::module), cgu_name);
+
+        if !self.available_cgus.contains(&cgu_name) {
+            let cgu_names: Vec<&str> =
+                self.available_cgus.items().map(|cgu| cgu.as_str()).into_sorted_stable_ord();
+            self.tcx.sess.emit_err(errors::NoModuleNamed {
+                span: attr.span,
+                user_path,
+                cgu_name,
+                cgu_names: cgu_names.join(", "),
+            });
+        }
+
+        self.cgu_reuse_tracker.set_expectation(
+            cgu_name,
+            &user_path,
+            attr.span,
+            expected_reuse,
+            comp_kind,
+        );
+    }
+
+    fn field(&self, attr: &ast::Attribute, name: Symbol) -> Symbol {
+        for item in attr.meta_item_list().unwrap_or_else(ThinVec::new) {
+            if item.has_name(name) {
+                if let Some(value) = item.value_str() {
+                    return value;
+                } else {
+                    self.tcx.sess.emit_fatal(errors::FieldAssociatedValueExpected {
+                        span: item.span(),
+                        name,
+                    });
+                }
+            }
+        }
+
+        self.tcx.sess.emit_fatal(errors::NoField { span: attr.span, name });
+    }
+
+    /// Scan for a `cfg="foo"` attribute and check whether we have a
+    /// cfg flag called `foo`.
+    fn check_config(&self, attr: &ast::Attribute) -> bool {
+        let config = &self.tcx.sess.parse_sess.config;
+        let value = self.field(attr, sym::cfg);
+        debug!("check_config(config={:?}, value={:?})", config, value);
+        if config.iter().any(|&(name, _)| name == value) {
+            debug!("check_config: matched");
+            return true;
+        }
+        debug!("check_config: no match found");
+        false
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
+pub enum CguReuse {
+    No,
+    PreLto,
+    PostLto,
+}
+
+impl fmt::Display for CguReuse {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            CguReuse::No => write!(f, "No"),
+            CguReuse::PreLto => write!(f, "PreLto "),
+            CguReuse::PostLto => write!(f, "PostLto "),
+        }
+    }
+}
+
+impl IntoDiagnosticArg for CguReuse {
+    fn into_diagnostic_arg(self) -> DiagnosticArgValue<'static> {
+        DiagnosticArgValue::Str(Cow::Owned(self.to_string()))
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub enum ComparisonKind {
+    Exact,
+    AtLeast,
+}
+
+struct TrackerData {
+    actual_reuse: FxHashMap<String, CguReuse>,
+    expected_reuse: FxHashMap<String, (String, Span, CguReuse, ComparisonKind)>,
+}
+
+pub struct CguReuseTracker {
+    data: Option<TrackerData>,
+}
+
+impl CguReuseTracker {
+    fn new() -> CguReuseTracker {
+        let data =
+            TrackerData { actual_reuse: Default::default(), expected_reuse: Default::default() };
+
+        CguReuseTracker { data: Some(data) }
+    }
+
+    fn new_disabled() -> CguReuseTracker {
+        CguReuseTracker { data: None }
+    }
+
+    pub fn set_actual_reuse(&mut self, cgu_name: &str, kind: CguReuse) {
+        if let Some(data) = &mut self.data {
+            debug!("set_actual_reuse({cgu_name:?}, {kind:?})");
+
+            let prev_reuse = data.actual_reuse.insert(cgu_name.to_string(), kind);
+            assert!(prev_reuse.is_none());
+        }
+    }
+
+    fn set_expectation(
+        &mut self,
+        cgu_name: Symbol,
+        cgu_user_name: &str,
+        error_span: Span,
+        expected_reuse: CguReuse,
+        comparison_kind: ComparisonKind,
+    ) {
+        if let Some(data) = &mut self.data {
+            debug!("set_expectation({cgu_name:?}, {expected_reuse:?}, {comparison_kind:?})");
+
+            data.expected_reuse.insert(
+                cgu_name.to_string(),
+                (cgu_user_name.to_string(), error_span, expected_reuse, comparison_kind),
+            );
+        }
+    }
+
+    fn check_expected_reuse(&self, sess: &Session) {
+        if let Some(ref data) = self.data {
+            for (cgu_name, &(ref cgu_user_name, ref error_span, expected_reuse, comparison_kind)) in
+                &data.expected_reuse
+            {
+                if let Some(&actual_reuse) = data.actual_reuse.get(cgu_name) {
+                    let (error, at_least) = match comparison_kind {
+                        ComparisonKind::Exact => (expected_reuse != actual_reuse, false),
+                        ComparisonKind::AtLeast => (actual_reuse < expected_reuse, true),
+                    };
+
+                    if error {
+                        let at_least = if at_least { 1 } else { 0 };
+                        errors::IncorrectCguReuseType {
+                            span: *error_span,
+                            cgu_user_name,
+                            actual_reuse,
+                            expected_reuse,
+                            at_least,
+                        };
+                    }
+                } else {
+                    sess.emit_fatal(errors::CguNotRecorded { cgu_user_name, cgu_name });
+                }
+            }
+        }
+    }
+}