Auto merge of #80993 - Aaron1011:collect-set-tokens, r=petrochenkov

Set tokens on AST node in `collect_tokens`

A new `HasTokens` trait is introduced, which is used to move logic from
the callers of `collect_tokens` into the body of `collect_tokens`.

In addition to reducing duplication, this paves the way for PR #80689,
which needs to perform additional logic during token collection.

1 files changed, 66 insertions, 10 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index e933eb434a1..ff89982a4ef 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -925,16 +925,6 @@ impl Stmt {
         }
     }
 
-    pub fn set_tokens(&mut self, tokens: Option<LazyTokenStream>) {
-        match self.kind {
-            StmtKind::Local(ref mut local) => local.tokens = tokens,
-            StmtKind::Item(ref mut item) => item.tokens = tokens,
-            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => expr.tokens = tokens,
-            StmtKind::Empty => {}
-            StmtKind::MacCall(ref mut mac) => mac.tokens = tokens,
-        }
-    }
-
     pub fn has_trailing_semicolon(&self) -> bool {
         match &self.kind {
             StmtKind::Semi(_) => true,
@@ -2890,3 +2880,69 @@ impl TryFrom<ItemKind> for ForeignItemKind {
 }
 
 pub type ForeignItem = Item<ForeignItemKind>;
+
+pub trait HasTokens {
+    /// Called by `Parser::collect_tokens` to store the collected
+    /// tokens inside an AST node
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream);
+}
+
+impl<T: HasTokens + 'static> HasTokens for P<T> {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+        (**self).finalize_tokens(tokens);
+    }
+}
+
+impl<T: HasTokens> HasTokens for Option<T> {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+        if let Some(inner) = self {
+            inner.finalize_tokens(tokens);
+        }
+    }
+}
+
+impl HasTokens for Attribute {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+        match &mut self.kind {
+            AttrKind::Normal(_, attr_tokens) => {
+                if attr_tokens.is_none() {
+                    *attr_tokens = Some(tokens);
+                }
+            }
+            AttrKind::DocComment(..) => {
+                panic!("Called finalize_tokens on doc comment attr {:?}", self)
+            }
+        }
+    }
+}
+
+impl HasTokens for Stmt {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+        let stmt_tokens = match self.kind {
+            StmtKind::Local(ref mut local) => &mut local.tokens,
+            StmtKind::Item(ref mut item) => &mut item.tokens,
+            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => &mut expr.tokens,
+            StmtKind::Empty => return,
+            StmtKind::MacCall(ref mut mac) => &mut mac.tokens,
+        };
+        if stmt_tokens.is_none() {
+            *stmt_tokens = Some(tokens);
+        }
+    }
+}
+
+macro_rules! derive_has_tokens {
+    ($($ty:path),*) => { $(
+        impl HasTokens for $ty {
+            fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+                if self.tokens.is_none() {
+                    self.tokens = Some(tokens);
+                }
+            }
+        }
+    )* }
+}
+
+derive_has_tokens! {
+    Item, Expr, Ty, AttrItem, Visibility, Path, Block, Pat
+}