iok_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{Fields, Item, parse_macro_input};
4
5#[proc_macro_attribute]
6pub fn ast_node(_attr: TokenStream, item: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(item as Item);
8
9    let expanded = match input {
10        // ---------------- STRUCTS ----------------
11        Item::Struct(mut struct_item) => {
12            let name = &struct_item.ident;
13
14            // Only support named fields
15            let fields = match &mut struct_item.fields {
16                Fields::Named(fields) => fields,
17                _ => panic!("`#[ast_node]` can only be applied to structs with named fields"),
18            };
19
20            // Add `span` if missing
21            let has_span = fields
22                .named
23                .iter()
24                .any(|f| f.ident.as_ref().map_or(false, |ident| ident == "span"));
25            if !has_span {
26                let span_field: syn::Field = syn::parse_quote! {
27                    pub span: core::ops::Range<usize>
28                };
29                fields.named.push(span_field);
30            }
31
32            quote! {
33                #struct_item
34
35                impl iok::Span for #name {
36                    fn span(&self) -> core::ops::Range<usize> {
37                        self.span.clone()
38                    }
39                }
40            }
41        }
42
43        // ---------------- ENUMS ----------------
44        Item::Enum(mut enum_item) => {
45            let name = &enum_item.ident;
46            let attrs = &enum_item.attrs;
47
48            let mut new_structs = Vec::new();
49
50            // Process variants
51            let variants = enum_item
52                .variants
53                .iter_mut()
54                .map(|variant| {
55                    let variant_name = &variant.ident;
56                    match &variant.fields {
57                        // Tuple variant with one field → use inner span
58                        Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
59                            quote! {
60                                #name::#variant_name(inner) => inner.span()
61                            }
62                        }
63
64                        // Unit variant → convert to wrapped struct
65                        Fields::Unit => {
66                            let struct_name = format_ident!("{}{}", name, variant_name);
67
68                            // Replace variant with tuple variant wrapping the struct
69                            variant.fields = syn::Fields::Unnamed(syn::parse_quote!((#struct_name)));
70
71                            // Create the span struct for this variant
72                            let new_struct: proc_macro2::TokenStream = quote! {
73                                #(#attrs)*
74                                pub struct #struct_name {
75                                    pub span: core::ops::Range<usize>
76                                }
77
78                                impl iok::Span for #struct_name {
79                                    fn span(&self) -> core::ops::Range<usize> {
80                                        self.span.clone()
81                                    }
82                                }
83
84                                impl From<core::ops::Range<usize>> for #struct_name {
85                                    fn from(span: core::ops::Range<usize>) -> Self {
86                                        #struct_name { span }
87                                    }
88                                }
89                            };
90                            new_structs.push(new_struct);
91
92                            quote! {
93                                #name::#variant_name(inner) => inner.span()
94                            }
95                        }
96
97                        // Named fields not supported
98                        Fields::Named(_) => {
99                            panic!("`#[ast_node]` for enums only supports unit or single-field tuple variants");
100                        }
101
102                        _ => panic!("Unexpected variant format"),
103                    }
104                })
105                .collect::<Vec<_>>();
106
107            quote! {
108                #enum_item
109                #(#new_structs)*
110
111                impl iok::Span for #name {
112                    fn span(&self) -> core::ops::Range<usize> {
113                        match self {
114                            #(#variants,)*
115                        }
116                    }
117                }
118            }
119        }
120
121        _ => panic!("`#[ast_node]` can only be applied to structs or enums"),
122    };
123
124    TokenStream::from(expanded)
125}
iok_macros/lib.rs

iok_macros/
lib.rs
