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