rustc_ast/
ast_traits.rs

1//! A set of traits implemented for various AST nodes,
2//! typically those used in AST fragments during macro expansion.
3//! The traits are not implemented exhaustively, only when actually necessary.
4
5use std::fmt;
6use std::marker::PhantomData;
7
8use crate::ptr::P;
9use crate::token::Nonterminal;
10use crate::tokenstream::LazyAttrTokenStream;
11use crate::{
12    Arm, AssocItem, AttrItem, AttrKind, AttrVec, Attribute, Block, Crate, Expr, ExprField,
13    FieldDef, ForeignItem, GenericParam, Item, NodeId, Param, Pat, PatField, Path, Stmt, StmtKind,
14    Ty, Variant, Visibility,
15};
16
17/// A utility trait to reduce boilerplate.
18/// Standard `Deref(Mut)` cannot be reused due to coherence.
19pub trait AstDeref {
20    type Target;
21    fn ast_deref(&self) -> &Self::Target;
22    fn ast_deref_mut(&mut self) -> &mut Self::Target;
23}
24
25macro_rules! impl_not_ast_deref {
26    ($($T:ty),+ $(,)?) => {
27        $(
28            impl !AstDeref for $T {}
29        )+
30    };
31}
32
33impl_not_ast_deref!(AssocItem, Expr, ForeignItem, Item, Stmt);
34
35impl<T> AstDeref for P<T> {
36    type Target = T;
37    fn ast_deref(&self) -> &Self::Target {
38        self
39    }
40    fn ast_deref_mut(&mut self) -> &mut Self::Target {
41        self
42    }
43}
44
45/// A trait for AST nodes having an ID.
46pub trait HasNodeId {
47    fn node_id(&self) -> NodeId;
48    fn node_id_mut(&mut self) -> &mut NodeId;
49}
50
51macro_rules! impl_has_node_id {
52    ($($T:ty),+ $(,)?) => {
53        $(
54            impl HasNodeId for $T {
55                fn node_id(&self) -> NodeId {
56                    self.id
57                }
58                fn node_id_mut(&mut self) -> &mut NodeId {
59                    &mut self.id
60                }
61            }
62        )+
63    };
64}
65
66impl_has_node_id!(
67    Arm,
68    AssocItem,
69    Crate,
70    Expr,
71    ExprField,
72    FieldDef,
73    ForeignItem,
74    GenericParam,
75    Item,
76    Param,
77    Pat,
78    PatField,
79    Stmt,
80    Ty,
81    Variant,
82);
83
84impl<T: AstDeref<Target: HasNodeId>> HasNodeId for T {
85    fn node_id(&self) -> NodeId {
86        self.ast_deref().node_id()
87    }
88    fn node_id_mut(&mut self) -> &mut NodeId {
89        self.ast_deref_mut().node_id_mut()
90    }
91}
92
93/// A trait for AST nodes having (or not having) collected tokens.
94pub trait HasTokens {
95    fn tokens(&self) -> Option<&LazyAttrTokenStream>;
96    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>>;
97}
98
99macro_rules! impl_has_tokens {
100    ($($T:ty),+ $(,)?) => {
101        $(
102            impl HasTokens for $T {
103                fn tokens(&self) -> Option<&LazyAttrTokenStream> {
104                    self.tokens.as_ref()
105                }
106                fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
107                    Some(&mut self.tokens)
108                }
109            }
110        )+
111    };
112}
113
114macro_rules! impl_has_tokens_none {
115    ($($T:ty),+ $(,)?) => {
116        $(
117            impl HasTokens for $T {
118                fn tokens(&self) -> Option<&LazyAttrTokenStream> {
119                    None
120                }
121                fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
122                    None
123                }
124            }
125        )+
126    };
127}
128
129impl_has_tokens!(AssocItem, AttrItem, Block, Expr, ForeignItem, Item, Pat, Path, Ty, Visibility);
130impl_has_tokens_none!(Arm, ExprField, FieldDef, GenericParam, Param, PatField, Variant);
131
132impl<T: AstDeref<Target: HasTokens>> HasTokens for T {
133    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
134        self.ast_deref().tokens()
135    }
136    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
137        self.ast_deref_mut().tokens_mut()
138    }
139}
140
141impl<T: HasTokens> HasTokens for Option<T> {
142    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
143        self.as_ref().and_then(|inner| inner.tokens())
144    }
145    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
146        self.as_mut().and_then(|inner| inner.tokens_mut())
147    }
148}
149
150impl HasTokens for StmtKind {
151    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
152        match self {
153            StmtKind::Let(local) => local.tokens.as_ref(),
154            StmtKind::Item(item) => item.tokens(),
155            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.tokens(),
156            StmtKind::Empty => None,
157            StmtKind::MacCall(mac) => mac.tokens.as_ref(),
158        }
159    }
160    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
161        match self {
162            StmtKind::Let(local) => Some(&mut local.tokens),
163            StmtKind::Item(item) => item.tokens_mut(),
164            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.tokens_mut(),
165            StmtKind::Empty => None,
166            StmtKind::MacCall(mac) => Some(&mut mac.tokens),
167        }
168    }
169}
170
171impl HasTokens for Stmt {
172    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
173        self.kind.tokens()
174    }
175    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
176        self.kind.tokens_mut()
177    }
178}
179
180impl HasTokens for Attribute {
181    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
182        match &self.kind {
183            AttrKind::Normal(normal) => normal.tokens.as_ref(),
184            kind @ AttrKind::DocComment(..) => {
185                panic!("Called tokens on doc comment attr {kind:?}")
186            }
187        }
188    }
189    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
190        Some(match &mut self.kind {
191            AttrKind::Normal(normal) => &mut normal.tokens,
192            kind @ AttrKind::DocComment(..) => {
193                panic!("Called tokens_mut on doc comment attr {kind:?}")
194            }
195        })
196    }
197}
198
199impl HasTokens for Nonterminal {
200    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
201        match self {
202            Nonterminal::NtItem(item) => item.tokens(),
203            Nonterminal::NtStmt(stmt) => stmt.tokens(),
204            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.tokens(),
205            Nonterminal::NtPat(pat) => pat.tokens(),
206            Nonterminal::NtTy(ty) => ty.tokens(),
207            Nonterminal::NtMeta(attr_item) => attr_item.tokens(),
208            Nonterminal::NtPath(path) => path.tokens(),
209            Nonterminal::NtVis(vis) => vis.tokens(),
210            Nonterminal::NtBlock(block) => block.tokens(),
211        }
212    }
213    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
214        match self {
215            Nonterminal::NtItem(item) => item.tokens_mut(),
216            Nonterminal::NtStmt(stmt) => stmt.tokens_mut(),
217            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.tokens_mut(),
218            Nonterminal::NtPat(pat) => pat.tokens_mut(),
219            Nonterminal::NtTy(ty) => ty.tokens_mut(),
220            Nonterminal::NtMeta(attr_item) => attr_item.tokens_mut(),
221            Nonterminal::NtPath(path) => path.tokens_mut(),
222            Nonterminal::NtVis(vis) => vis.tokens_mut(),
223            Nonterminal::NtBlock(block) => block.tokens_mut(),
224        }
225    }
226}
227
228/// A trait for AST nodes having (or not having) attributes.
229pub trait HasAttrs {
230    /// This is `true` if this `HasAttrs` might support 'custom' (proc-macro) inner
231    /// attributes. Attributes like `#![cfg]` and `#![cfg_attr]` are not
232    /// considered 'custom' attributes.
233    ///
234    /// If this is `false`, then this `HasAttrs` definitely does
235    /// not support 'custom' inner attributes, which enables some optimizations
236    /// during token collection.
237    const SUPPORTS_CUSTOM_INNER_ATTRS: bool;
238    fn attrs(&self) -> &[Attribute];
239    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec));
240}
241
242macro_rules! impl_has_attrs {
243    (const SUPPORTS_CUSTOM_INNER_ATTRS: bool = $inner:literal, $($T:ty),+ $(,)?) => {
244        $(
245            impl HasAttrs for $T {
246                const SUPPORTS_CUSTOM_INNER_ATTRS: bool = $inner;
247
248                #[inline]
249                fn attrs(&self) -> &[Attribute] {
250                    &self.attrs
251                }
252
253                fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
254                    f(&mut self.attrs)
255                }
256            }
257        )+
258    };
259}
260
261macro_rules! impl_has_attrs_none {
262    ($($T:ty),+ $(,)?) => {
263        $(
264            impl HasAttrs for $T {
265                const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false;
266                fn attrs(&self) -> &[Attribute] {
267                    &[]
268                }
269                fn visit_attrs(&mut self, _f: impl FnOnce(&mut AttrVec)) {}
270            }
271        )+
272    };
273}
274
275impl_has_attrs!(
276    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true,
277    AssocItem,
278    ForeignItem,
279    Item,
280);
281impl_has_attrs!(
282    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false,
283    Arm,
284    Crate,
285    Expr,
286    ExprField,
287    FieldDef,
288    GenericParam,
289    Param,
290    PatField,
291    Variant,
292);
293impl_has_attrs_none!(Attribute, AttrItem, Block, Pat, Path, Ty, Visibility);
294
295impl<T: AstDeref<Target: HasAttrs>> HasAttrs for T {
296    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::Target::SUPPORTS_CUSTOM_INNER_ATTRS;
297    fn attrs(&self) -> &[Attribute] {
298        self.ast_deref().attrs()
299    }
300    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
301        self.ast_deref_mut().visit_attrs(f)
302    }
303}
304
305impl<T: HasAttrs> HasAttrs for Option<T> {
306    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::SUPPORTS_CUSTOM_INNER_ATTRS;
307    fn attrs(&self) -> &[Attribute] {
308        self.as_ref().map(|inner| inner.attrs()).unwrap_or(&[])
309    }
310    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
311        if let Some(inner) = self.as_mut() {
312            inner.visit_attrs(f);
313        }
314    }
315}
316
317impl HasAttrs for StmtKind {
318    // This might be a `StmtKind::Item`, which contains
319    // an item that supports inner attrs.
320    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true;
321
322    fn attrs(&self) -> &[Attribute] {
323        match self {
324            StmtKind::Let(local) => &local.attrs,
325            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.attrs(),
326            StmtKind::Item(item) => item.attrs(),
327            StmtKind::Empty => &[],
328            StmtKind::MacCall(mac) => &mac.attrs,
329        }
330    }
331
332    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
333        match self {
334            StmtKind::Let(local) => f(&mut local.attrs),
335            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.visit_attrs(f),
336            StmtKind::Item(item) => item.visit_attrs(f),
337            StmtKind::Empty => {}
338            StmtKind::MacCall(mac) => f(&mut mac.attrs),
339        }
340    }
341}
342
343impl HasAttrs for Stmt {
344    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = StmtKind::SUPPORTS_CUSTOM_INNER_ATTRS;
345    fn attrs(&self) -> &[Attribute] {
346        self.kind.attrs()
347    }
348    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
349        self.kind.visit_attrs(f);
350    }
351}
352
353/// A newtype around an AST node that implements the traits above if the node implements them.
354pub struct AstNodeWrapper<Wrapped, Tag> {
355    pub wrapped: Wrapped,
356    pub tag: PhantomData<Tag>,
357}
358
359impl<Wrapped, Tag> AstNodeWrapper<Wrapped, Tag> {
360    pub fn new(wrapped: Wrapped, _tag: Tag) -> AstNodeWrapper<Wrapped, Tag> {
361        AstNodeWrapper { wrapped, tag: Default::default() }
362    }
363}
364
365impl<Wrapped, Tag> AstDeref for AstNodeWrapper<Wrapped, Tag> {
366    type Target = Wrapped;
367    fn ast_deref(&self) -> &Self::Target {
368        &self.wrapped
369    }
370    fn ast_deref_mut(&mut self) -> &mut Self::Target {
371        &mut self.wrapped
372    }
373}
374
375impl<Wrapped: fmt::Debug, Tag> fmt::Debug for AstNodeWrapper<Wrapped, Tag> {
376    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
377        f.debug_struct("AstNodeWrapper")
378            .field("wrapped", &self.wrapped)
379            .field("tag", &self.tag)
380            .finish()
381    }
382}