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

rustc_ast/
ast_traits.rs
