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