Skip to main content

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