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