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, WherePredicate,
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    WherePredicate,
83);
84
85impl<T: AstDeref<Target: HasNodeId>> HasNodeId for T {
86    fn node_id(&self) -> NodeId {
87        self.ast_deref().node_id()
88    }
89    fn node_id_mut(&mut self) -> &mut NodeId {
90        self.ast_deref_mut().node_id_mut()
91    }
92}
93
94/// A trait for AST nodes having (or not having) collected tokens.
95pub trait HasTokens {
96    fn tokens(&self) -> Option<&LazyAttrTokenStream>;
97    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>>;
98}
99
100macro_rules! impl_has_tokens {
101    ($($T:ty),+ $(,)?) => {
102        $(
103            impl HasTokens for $T {
104                fn tokens(&self) -> Option<&LazyAttrTokenStream> {
105                    self.tokens.as_ref()
106                }
107                fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
108                    Some(&mut self.tokens)
109                }
110            }
111        )+
112    };
113}
114
115macro_rules! impl_has_tokens_none {
116    ($($T:ty),+ $(,)?) => {
117        $(
118            impl HasTokens for $T {
119                fn tokens(&self) -> Option<&LazyAttrTokenStream> {
120                    None
121                }
122                fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
123                    None
124                }
125            }
126        )+
127    };
128}
129
130impl_has_tokens!(AssocItem, AttrItem, Block, Expr, ForeignItem, Item, Pat, Path, Ty, Visibility);
131impl_has_tokens_none!(
132    Arm,
133    ExprField,
134    FieldDef,
135    GenericParam,
136    Param,
137    PatField,
138    Variant,
139    WherePredicate
140);
141
142impl<T: AstDeref<Target: HasTokens>> HasTokens for T {
143    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
144        self.ast_deref().tokens()
145    }
146    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
147        self.ast_deref_mut().tokens_mut()
148    }
149}
150
151impl<T: HasTokens> HasTokens for Option<T> {
152    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
153        self.as_ref().and_then(|inner| inner.tokens())
154    }
155    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
156        self.as_mut().and_then(|inner| inner.tokens_mut())
157    }
158}
159
160impl HasTokens for StmtKind {
161    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
162        match self {
163            StmtKind::Let(local) => local.tokens.as_ref(),
164            StmtKind::Item(item) => item.tokens(),
165            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.tokens(),
166            StmtKind::Empty => None,
167            StmtKind::MacCall(mac) => mac.tokens.as_ref(),
168        }
169    }
170    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
171        match self {
172            StmtKind::Let(local) => Some(&mut local.tokens),
173            StmtKind::Item(item) => item.tokens_mut(),
174            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.tokens_mut(),
175            StmtKind::Empty => None,
176            StmtKind::MacCall(mac) => Some(&mut mac.tokens),
177        }
178    }
179}
180
181impl HasTokens for Stmt {
182    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
183        self.kind.tokens()
184    }
185    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
186        self.kind.tokens_mut()
187    }
188}
189
190impl HasTokens for Attribute {
191    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
192        match &self.kind {
193            AttrKind::Normal(normal) => normal.tokens.as_ref(),
194            kind @ AttrKind::DocComment(..) => {
195                panic!("Called tokens on doc comment attr {kind:?}")
196            }
197        }
198    }
199    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
200        Some(match &mut self.kind {
201            AttrKind::Normal(normal) => &mut normal.tokens,
202            kind @ AttrKind::DocComment(..) => {
203                panic!("Called tokens_mut on doc comment attr {kind:?}")
204            }
205        })
206    }
207}
208
209impl HasTokens for Nonterminal {
210    fn tokens(&self) -> Option<&LazyAttrTokenStream> {
211        match self {
212            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.tokens(),
213            Nonterminal::NtBlock(block) => block.tokens(),
214        }
215    }
216    fn tokens_mut(&mut self) -> Option<&mut Option<LazyAttrTokenStream>> {
217        match self {
218            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.tokens_mut(),
219            Nonterminal::NtBlock(block) => block.tokens_mut(),
220        }
221    }
222}
223
224/// A trait for AST nodes having (or not having) attributes.
225pub trait HasAttrs {
226    /// This is `true` if this `HasAttrs` might support 'custom' (proc-macro) inner
227    /// attributes. Attributes like `#![cfg]` and `#![cfg_attr]` are not
228    /// considered 'custom' attributes.
229    ///
230    /// If this is `false`, then this `HasAttrs` definitely does
231    /// not support 'custom' inner attributes, which enables some optimizations
232    /// during token collection.
233    const SUPPORTS_CUSTOM_INNER_ATTRS: bool;
234    fn attrs(&self) -> &[Attribute];
235    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec));
236}
237
238macro_rules! impl_has_attrs {
239    (const SUPPORTS_CUSTOM_INNER_ATTRS: bool = $inner:literal, $($T:ty),+ $(,)?) => {
240        $(
241            impl HasAttrs for $T {
242                const SUPPORTS_CUSTOM_INNER_ATTRS: bool = $inner;
243
244                #[inline]
245                fn attrs(&self) -> &[Attribute] {
246                    &self.attrs
247                }
248
249                fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
250                    f(&mut self.attrs)
251                }
252            }
253        )+
254    };
255}
256
257macro_rules! impl_has_attrs_none {
258    ($($T:ty),+ $(,)?) => {
259        $(
260            impl HasAttrs for $T {
261                const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false;
262                fn attrs(&self) -> &[Attribute] {
263                    &[]
264                }
265                fn visit_attrs(&mut self, _f: impl FnOnce(&mut AttrVec)) {}
266            }
267        )+
268    };
269}
270
271impl_has_attrs!(
272    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true,
273    AssocItem,
274    ForeignItem,
275    Item,
276);
277impl_has_attrs!(
278    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false,
279    Arm,
280    Crate,
281    Expr,
282    ExprField,
283    FieldDef,
284    GenericParam,
285    Param,
286    PatField,
287    Variant,
288    WherePredicate,
289);
290impl_has_attrs_none!(Attribute, AttrItem, Block, Pat, Path, Ty, Visibility);
291
292impl<T: AstDeref<Target: HasAttrs>> HasAttrs for T {
293    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::Target::SUPPORTS_CUSTOM_INNER_ATTRS;
294    fn attrs(&self) -> &[Attribute] {
295        self.ast_deref().attrs()
296    }
297    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
298        self.ast_deref_mut().visit_attrs(f)
299    }
300}
301
302impl<T: HasAttrs> HasAttrs for Option<T> {
303    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::SUPPORTS_CUSTOM_INNER_ATTRS;
304    fn attrs(&self) -> &[Attribute] {
305        self.as_ref().map(|inner| inner.attrs()).unwrap_or(&[])
306    }
307    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
308        if let Some(inner) = self.as_mut() {
309            inner.visit_attrs(f);
310        }
311    }
312}
313
314impl HasAttrs for StmtKind {
315    // This might be a `StmtKind::Item`, which contains
316    // an item that supports inner attrs.
317    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true;
318
319    fn attrs(&self) -> &[Attribute] {
320        match self {
321            StmtKind::Let(local) => &local.attrs,
322            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.attrs(),
323            StmtKind::Item(item) => item.attrs(),
324            StmtKind::Empty => &[],
325            StmtKind::MacCall(mac) => &mac.attrs,
326        }
327    }
328
329    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
330        match self {
331            StmtKind::Let(local) => f(&mut local.attrs),
332            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.visit_attrs(f),
333            StmtKind::Item(item) => item.visit_attrs(f),
334            StmtKind::Empty => {}
335            StmtKind::MacCall(mac) => f(&mut mac.attrs),
336        }
337    }
338}
339
340impl HasAttrs for Stmt {
341    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = StmtKind::SUPPORTS_CUSTOM_INNER_ATTRS;
342    fn attrs(&self) -> &[Attribute] {
343        self.kind.attrs()
344    }
345    fn visit_attrs(&mut self, f: impl FnOnce(&mut AttrVec)) {
346        self.kind.visit_attrs(f);
347    }
348}
349
350/// A newtype around an AST node that implements the traits above if the node implements them.
351pub struct AstNodeWrapper<Wrapped, Tag> {
352    pub wrapped: Wrapped,
353    pub tag: PhantomData<Tag>,
354}
355
356impl<Wrapped, Tag> AstNodeWrapper<Wrapped, Tag> {
357    pub fn new(wrapped: Wrapped, _tag: Tag) -> AstNodeWrapper<Wrapped, Tag> {
358        AstNodeWrapper { wrapped, tag: Default::default() }
359    }
360}
361
362impl<Wrapped, Tag> AstDeref for AstNodeWrapper<Wrapped, Tag> {
363    type Target = Wrapped;
364    fn ast_deref(&self) -> &Self::Target {
365        &self.wrapped
366    }
367    fn ast_deref_mut(&mut self) -> &mut Self::Target {
368        &mut self.wrapped
369    }
370}
371
372impl<Wrapped: fmt::Debug, Tag> fmt::Debug for AstNodeWrapper<Wrapped, Tag> {
373    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
374        f.debug_struct("AstNodeWrapper")
375            .field("wrapped", &self.wrapped)
376            .field("tag", &self.tag)
377            .finish()
378    }
379}