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        if let Statement::Declaration(stmt) = self {
51            Some(stmt.ident.to_string())
52        } else {
53            None
54        }
55    }
56}
57
58impl<T: NamedNode + TokRepr> TokRepr for Spanned<T> {
59    fn tok_repr(&self) -> TokenStream {
60        let node = self.node().tok_repr();
61        let span = self.span().tok_repr();
62
63        if let Some(name) = self.name()
64            && let Some(suffix) = name.strip_prefix("#")
65        {
66            let ident = format_ident!("{}", suffix);
67
68            return quote! {
69                Spanned::new(#ident.to_owned().into(), #span)
70            };
71        }
72
73        quote! {
74            Spanned::new(#node, #span)
75        }
76    }
77}
78
79impl TokRepr for Ident {
80    fn tok_repr(&self) -> TokenStream {
81        let name = self.name();
82        if let Some(name) = name.strip_prefix("#") {
83            let ident = format_ident!("{}", name);
84            quote! {
85                Ident::from(#ident.to_owned())
86            }
87        } else {
88            let name = name.as_str();
89            quote! {
90                Ident::new(#name.to_string())
91            }
92        }
93    }
94}