Skip to main content

wgsl_parse/
tokrepr.rs

1//! Turn an instance into a `TokenStream` that represents the instance.
2
3use tokrepr::TokRepr;
4use tokrepr::proc_macro2::TokenStream;
5use tokrepr::quote::{format_ident, quote};
6
7use crate::{SyntaxNode, span::Spanned, syntax::*};
8
9/// named nodes are those that have an identifier. We use this trait to conditionally
10/// implement TokRepr for Spanned<NamedNode>, which allows code injection in `quote_*!`
11/// macros from the `wesl` crate.
12trait NamedNode {
13    fn name(&self) -> Option<String>;
14}
15
16impl NamedNode for GlobalDeclaration {
17    fn name(&self) -> Option<String> {
18        self.ident().map(|id| id.to_string())
19    }
20}
21
22impl NamedNode for StructMember {
23    fn name(&self) -> Option<String> {
24        Some(self.ident.to_string())
25    }
26}
27
28impl NamedNode for Attribute {
29    fn name(&self) -> Option<String> {
30        if let Attribute::Custom(attr) = self {
31            Some(attr.name.to_string())
32        } else {
33            None
34        }
35    }
36}
37
38impl NamedNode for Expression {
39    fn name(&self) -> Option<String> {
40        if let Expression::TypeOrIdentifier(ty) = self {
41            Some(ty.ident.to_string())
42        } else {
43            None
44        }
45    }
46}
47
48impl NamedNode for Statement {
49    fn name(&self) -> Option<String> {
50        // COMBAK: this nesting is hell. Hopefully if-let chains will stabilize soon.
51        if let Statement::Compound(stmt) = self
52            && stmt.statements.is_empty()
53            && let [attr] = stmt.attributes.as_slice()
54            && let Attribute::Custom(attr) = attr.node()
55        {
56            return Some(attr.name.to_string());
57        }
58        None
59    }
60}
61
62impl<T: NamedNode + TokRepr> TokRepr for Spanned<T> {
63    fn tok_repr(&self) -> TokenStream {
64        let node = self.node().tok_repr();
65        let span = self.span().tok_repr();
66
67        if let Some(name) = self.name()
68            && let Some(suffix) = name.strip_prefix("#")
69        {
70            let ident = format_ident!("{}", suffix);
71
72            return quote! {
73                Spanned::new(#ident.to_owned().into(), #span)
74            };
75        }
76
77        quote! {
78            Spanned::new(#node, #span)
79        }
80    }
81}
82
83impl TokRepr for Ident {
84    fn tok_repr(&self) -> TokenStream {
85        let name = self.name();
86        if let Some(name) = name.strip_prefix("#") {
87            let ident = format_ident!("{}", name);
88            quote! {
89                Ident::from(#ident.to_owned())
90            }
91        } else {
92            let name = name.as_str();
93            quote! {
94                Ident::new(#name.to_string())
95            }
96        }
97    }
98}