Lower impl restriction information to TraitDef

compiler/rustc_hir_analysis/src/collect.rs

@@ -893,11 +893,25 @@ fn adt_def(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::AdtDef<'_> {
 fn trait_def(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::TraitDef {
     let item = tcx.hir_expect_item(def_id);
 
-    let (constness, is_alias, is_auto, safety) = match item.kind {
-        hir::ItemKind::Trait(constness, is_auto, safety, ..) => {
-            (constness, false, is_auto == hir::IsAuto::Yes, safety)
-        }
-        hir::ItemKind::TraitAlias(constness, ..) => (constness, true, false, hir::Safety::Safe),
+    let (constness, is_alias, is_auto, safety, impl_restriction) = match item.kind {
+        hir::ItemKind::Trait(constness, is_auto, safety, impl_restriction, ..) => (
+            constness,
+            false,
+            is_auto == hir::IsAuto::Yes,
+            safety,
+            if let hir::RestrictionKind::Restricted { path, shorthand: _ } = impl_restriction.kind {
+                ty::trait_def::ImplRestrictionKind::Restricted(path.res, impl_restriction.span)
+            } else {
+                ty::trait_def::ImplRestrictionKind::Unrestricted
+            },
+        ),
+        hir::ItemKind::TraitAlias(constness, ..) => (
+            constness,
+            true,
+            false,
+            hir::Safety::Safe,
+            ty::trait_def::ImplRestrictionKind::Unrestricted,
+        ),
         _ => span_bug!(item.span, "trait_def_of_item invoked on non-trait"),
     };
 
@@ -946,6 +960,7 @@ fn trait_def(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::TraitDef {
         def_id: def_id.to_def_id(),
         safety,
         constness,
+        impl_restriction,
         paren_sugar,
         has_auto_impl: is_auto,
         is_marker,

compiler/rustc_middle/src/ty/trait_def.rs

@@ -24,6 +24,9 @@ pub struct TraitDef {
     /// Whether this trait is `const`.
     pub constness: hir::Constness,
 
+    /// Restrictions on trait implementations.
+    pub impl_restriction: ImplRestrictionKind,
+
     /// If `true`, then this trait had the `#[rustc_paren_sugar]`
     /// attribute, indicating that it should be used with `Foo()`
     /// sugar. This is a temporary thing -- eventually any trait will
@@ -97,6 +100,52 @@ pub enum TraitSpecializationKind {
     AlwaysApplicable,
 }
 
+/// Whether the trait implementation is unrestricted or restricted within a specific module.
+#[derive(HashStable, PartialEq, Clone, Copy, Encodable, Decodable)]
+pub enum ImplRestrictionKind {
+    /// The restriction does not affect this trait, and it can be implemented anywhere.
+    Unrestricted,
+    /// This trait can only be implemented within the specified module.
+    Restricted(DefId, Span),
+}
+
+impl ImplRestrictionKind {
+    /// Returns `true` if the behavior is allowed/unrestricted in the given module.
+    /// A value of `false` indicates that the behavior is prohibited.
+    pub fn is_allowed_in(self, module: DefId, tcx: TyCtxt<'_>) -> bool {
+        match self {
+            ImplRestrictionKind::Unrestricted => true,
+            ImplRestrictionKind::Restricted(restricted_to, _) => {
+                tcx.is_descendant_of(module, restricted_to)
+            }
+        }
+    }
+
+    /// Obtain the [`Span`] of the restriction. Panics if the restriction is unrestricted.
+    pub fn expect_span(self) -> Span {
+        match self {
+            ImplRestrictionKind::Unrestricted => {
+                bug!("called `expect_span` on an unrestricted item")
+            }
+            ImplRestrictionKind::Restricted(_, span) => span,
+        }
+    }
+
+    /// Obtain the path of the restriction. If unrestricted, an empty string is returned.
+    pub fn restriction_path(self, tcx: TyCtxt<'_>, krate: rustc_span::def_id::CrateNum) -> String {
+        match self {
+            ImplRestrictionKind::Unrestricted => String::new(),
+            ImplRestrictionKind::Restricted(restricted_to, _) => {
+                if restricted_to.krate == krate {
+                    tcx.def_path_str(restricted_to)
+                } else {
+                    tcx.crate_name(restricted_to.krate).to_ident_string()
+                }
+            }
+        }
+    }
+}
+
 #[derive(Default, Debug, HashStable)]
 pub struct TraitImpls {
     blanket_impls: Vec<DefId>,