Introduce restrictions to the AST

Author: jhpratt

compiler/rustc_ast/src/ast.rs

@@ -32,7 +32,7 @@ use rustc_data_structures::tagged_ptr::Tag;
 use rustc_macros::{Decodable, Encodable, HashStable_Generic};
 pub use rustc_span::AttrId;
 use rustc_span::source_map::{Spanned, respan};
-use rustc_span::{ErrorGuaranteed, Ident, Span, Symbol, kw, sym};
+use rustc_span::{DUMMY_SP, ErrorGuaranteed, Ident, Span, Symbol, kw, sym};
 use thin_vec::{ThinVec, thin_vec};
 
 pub use crate::format::*;
@@ -3316,6 +3316,42 @@ impl VisibilityKind {
     }
 }
 
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub struct Restriction {
+    pub kind: RestrictionKind,
+    pub span: Span,
+    pub tokens: Option<LazyAttrTokenStream>,
+}
+
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub enum RestrictionKind {
+    Unrestricted,
+    Restricted { path: P<Path>, id: NodeId, shorthand: bool },
+    Implied,
+}
+
+impl Restriction {
+    pub fn unrestricted() -> Self {
+        Restriction { kind: RestrictionKind::Unrestricted, span: DUMMY_SP, tokens: None }
+    }
+
+    pub fn restricted(path: P<Path>, id: NodeId, shorthand: bool) -> Self {
+        Restriction {
+            kind: RestrictionKind::Restricted { path, id, shorthand },
+            span: DUMMY_SP,
+            tokens: None,
+        }
+    }
+
+    pub fn implied() -> Self {
+        Restriction { kind: RestrictionKind::Implied, span: DUMMY_SP, tokens: None }
+    }
+
+    pub fn with_span(self, span: Span) -> Self {
+        Restriction { span, ..self }
+    }
+}
+
 /// Field definition in a struct, variant or union.
 ///
 /// E.g., `bar: usize` as in `struct Foo { bar: usize }`.

compiler/rustc_ast/src/ast_traits.rs

@@ -9,8 +9,8 @@ use crate::ptr::P;
 use crate::tokenstream::LazyAttrTokenStream;
 use crate::{
     Arm, AssocItem, AttrItem, AttrKind, AttrVec, Attribute, Block, Crate, Expr, ExprField,
-    FieldDef, ForeignItem, GenericParam, Item, NodeId, Param, Pat, PatField, Path, Stmt, StmtKind,
-    Ty, Variant, Visibility, WherePredicate,
+    FieldDef, ForeignItem, GenericParam, Item, NodeId, Param, Pat, PatField, Path, Restriction,
+    Stmt, StmtKind, Ty, Variant, Visibility, WherePredicate,
 };
 
 /// A trait for AST nodes having an ID.
@@ -98,7 +98,19 @@ macro_rules! impl_has_tokens_none {
     };
 }
 
-impl_has_tokens!(AssocItem, AttrItem, Block, Expr, ForeignItem, Item, Pat, Path, Ty, Visibility);
+impl_has_tokens!(
+    AssocItem,
+    AttrItem,
+    Block,
+    Expr,
+    ForeignItem,
+    Item,
+    Pat,
+    Path,
+    Restriction,
+    Ty,
+    Visibility
+);
 impl_has_tokens_none!(
     Arm,
     ExprField,
@@ -243,7 +255,7 @@ impl_has_attrs!(
     Variant,
     WherePredicate,
 );
-impl_has_attrs_none!(Attribute, AttrItem, Block, Pat, Path, Ty, Visibility);
+impl_has_attrs_none!(Attribute, AttrItem, Block, Pat, Path, Restriction, Ty, Visibility);
 
 impl<T: HasAttrs> HasAttrs for P<T> {
     const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::SUPPORTS_CUSTOM_INNER_ATTRS;

compiler/rustc_ast/src/mut_visit.rs

@@ -327,6 +327,10 @@ pub trait MutVisitor: Sized {
         walk_vis(self, vis);
     }
 
+    fn visit_restriction(&mut self, restriction: &mut Restriction) {
+        walk_restriction(self, restriction);
+    }
+
     fn visit_id(&mut self, _id: &mut NodeId) {
         // Do nothing.
     }
@@ -1405,6 +1409,18 @@ fn walk_vis<T: MutVisitor>(vis: &mut T, visibility: &mut Visibility) {
     vis.visit_span(span);
 }
 
+fn walk_restriction<T: MutVisitor>(vis: &mut T, restriction: &mut Restriction) {
+    let Restriction { kind, span, tokens: _ } = restriction;
+    match kind {
+        RestrictionKind::Unrestricted | RestrictionKind::Implied => {}
+        RestrictionKind::Restricted { path, id, shorthand: _ } => {
+            vis.visit_id(id);
+            vis.visit_path(path);
+        }
+    }
+    vis.visit_span(span);
+}
+
 fn walk_capture_by<T: MutVisitor>(vis: &mut T, capture_by: &mut CaptureBy) {
     match capture_by {
         CaptureBy::Ref => {}

compiler/rustc_ast/src/visit.rs

@@ -268,6 +268,9 @@ pub trait Visitor<'ast>: Sized {
     fn visit_vis(&mut self, vis: &'ast Visibility) -> Self::Result {
         walk_vis(self, vis)
     }
+    fn visit_restriction(&mut self, restriction: &'ast Restriction) -> Self::Result {
+        walk_restriction(self, restriction)
+    }
     fn visit_fn_ret_ty(&mut self, ret_ty: &'ast FnRetTy) -> Self::Result {
         walk_fn_ret_ty(self, ret_ty)
     }
@@ -1516,6 +1519,20 @@ pub fn walk_vis<'a, V: Visitor<'a>>(visitor: &mut V, vis: &'a Visibility) -> V::
     V::Result::output()
 }
 
+pub fn walk_restriction<'a, V: Visitor<'a>>(
+    visitor: &mut V,
+    restriction: &'a Restriction,
+) -> V::Result {
+    let Restriction { kind, span: _, tokens: _ } = restriction;
+    match kind {
+        RestrictionKind::Unrestricted | RestrictionKind::Implied => {}
+        RestrictionKind::Restricted { path, id, shorthand: _ } => {
+            try_visit!(visitor.visit_path(path, *id));
+        }
+    }
+    V::Result::output()
+}
+
 pub fn walk_attribute<'a, V: Visitor<'a>>(visitor: &mut V, attr: &'a Attribute) -> V::Result {
     let Attribute { kind, id: _, style: _, span: _ } = attr;
     match kind {

compiler/rustc_ast_pretty/src/pprust/state.rs

@@ -1071,6 +1071,10 @@ pub trait PrintState<'a>: std::ops::Deref<Target = pp::Printer> + std::ops::Dere
         Self::to_string(|s| s.print_visibility(v))
     }
 
+    fn restriction_to_string(&self, kw: &'static str, r: &ast::Restriction) -> String {
+        Self::to_string(|s| s.print_restriction(kw, r))
+    }
+
     fn block_to_string(&self, blk: &ast::Block) -> String {
         Self::to_string(|s| {
             let (cb, ib) = s.head("");

compiler/rustc_ast_pretty/src/pprust/state/item.rs

@@ -473,6 +473,22 @@ impl<'a> State<'a> {
         }
     }
 
+    // FIXME(jhpratt) make `kw` into a const generic once permitted
+    pub(crate) fn print_restriction(&mut self, kw: &'static str, restriction: &ast::Restriction) {
+        match restriction.kind {
+            ast::RestrictionKind::Unrestricted => self.word_nbsp(kw),
+            ast::RestrictionKind::Restricted { ref path, shorthand, id: _ } => {
+                let path = Self::to_string(|s| s.print_path(path, false, 0));
+                if shorthand {
+                    self.word_nbsp(format!("{kw}({path})"))
+                } else {
+                    self.word_nbsp(format!("{kw}(in {path})"))
+                }
+            }
+            ast::RestrictionKind::Implied => {}
+        }
+    }
+
     fn print_defaultness(&mut self, defaultness: ast::Defaultness) {
         if let ast::Defaultness::Default(_) = defaultness {
             self.word_nbsp("default");

tests/ui/macros/stringify.rs

@@ -793,6 +793,16 @@ fn test_vis() {
     // Attributes are not allowed on visibilities.
 }
 
+#[test]
+fn test_impl_restriction() {
+    assert_eq!(stringify!(pub impl(crate) trait Foo {}), "pub impl(crate) trait Foo {}");
+    assert_eq!(stringify!(pub impl(in crate) trait Foo {}), "pub impl(in crate) trait Foo {}");
+    assert_eq!(stringify!(pub impl(super) trait Foo {}), "pub impl(super) trait Foo {}");
+    assert_eq!(stringify!(pub impl(in super) trait Foo {}), "pub impl(in super) trait Foo {}");
+    assert_eq!(stringify!(pub impl(self) trait Foo {}), "pub impl(self) trait Foo {}");
+    assert_eq!(stringify!(pub impl(in self) trait Foo {}), "pub impl(in self) trait Foo {}");
+}
+
 macro_rules! p {
     ([$($tt:tt)*], $s:literal) => {
         assert_eq!(stringify!($($tt)*), $s);