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

rustc_ast/
ast_traits.rs
