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