Skip to main content

tytro/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use proc_macro::TokenStream;
4use quote::{ToTokens, format_ident, quote, quote_spanned};
5use syn::{Data, Field, Fields, GenericParam, parse_macro_input, spanned::Spanned};
6
7#[proc_macro_attribute]
8pub fn tytro(_cfg: TokenStream, item: TokenStream) -> TokenStream {
9    let item = parse_macro_input!(item as syn::DeriveInput);
10
11    assert!(item.attrs.is_empty(), "attribute is unsupported");
12
13    let find_rec = |fields: &Fields| {
14        let mut found = None;
15        for (i, f) in fields.iter().enumerate() {
16            assert!(f.attrs.is_empty(), "attribute is unsupported");
17            match &f.ty {
18                syn::Type::Path(p) if p.path.is_ident("Self") => {}
19                _ => continue,
20            }
21            let old = found.replace(i);
22            assert!(
23                old.is_none(),
24                "Self can only appear at most once per variant"
25            );
26        }
27        found
28    };
29
30    let variants = match &item.data {
31        Data::Struct(data) => vec![(&item.ident, find_rec(&data.fields), &data.fields)],
32        Data::Enum(data) => Vec::from_iter(data.variants.iter().map(|v| {
33            assert!(v.discriminant.is_none(), "discriminant is unsupported");
34            assert!(v.attrs.is_empty(), "attribute is unsupported");
35            (&v.ident, find_rec(&v.fields), &v.fields)
36        })),
37        Data::Union(_) => panic!("union is unsupported"),
38    };
39
40    assert!(
41        variants.iter().any(|(_, rec, _)| rec.is_some()),
42        "Self must appear at least once"
43    );
44
45    let syn::DeriveInput {
46        vis,
47        ident,
48        generics:
49            syn::Generics {
50                params,
51                where_clause,
52                ..
53            },
54        ..
55    } = &item;
56
57    let ident_mod = format_ident!("__{}__mod_tytro__", ident);
58    let ident_ref = format_ident!("{}Ref", ident);
59    let ident_mut = format_ident!("{}Mut", ident);
60    let ident_f = format_ident!("{}F", ident);
61    let ident_f_ref = format_ident!("{}FRef", ident);
62    let ident_f_mut = format_ident!("{}FMut", ident);
63    let ident_data = format_ident!("__{}__data_tytro__", ident);
64
65    let where_clause = where_clause.iter().flat_map(|c| c.predicates.iter());
66    let where_clause = quote! {
67        where #(#where_clause,)*
68    };
69
70    let args = params.iter().map::<&dyn ToTokens, _>(|param| match param {
71        GenericParam::Lifetime(p) => &p.lifetime,
72        GenericParam::Type(p) => &p.ident,
73        GenericParam::Const(p) => &p.ident,
74    });
75    let args = quote! { #(#args),* };
76
77    let params_impl = params.iter().map(|param| {
78        let mut param = param.clone();
79        match &mut param {
80            GenericParam::Lifetime(_) => {}
81            GenericParam::Type(param) => param.default = None,
82            GenericParam::Const(param) => param.default = None,
83        }
84        param
85    });
86    let params_impl = quote! { #(#params_impl),* };
87
88    let f_data = |this, pre, is_struct| {
89        let where_clause = &where_clause;
90        variants.iter().map(move |(_, rec, f)| {
91            let fields = f.iter().enumerate();
92            let fields = fields.map(|(i, f)| {
93                let Field {
94                    vis,
95                    ident,
96                    colon_token,
97                    ty,
98                    ..
99                } = f;
100                if Some(i) == *rec {
101                    quote! { #vis #ident #colon_token #this }
102                } else {
103                    quote! { #vis #ident #colon_token #pre #ty }
104                }
105            });
106            match f {
107                Fields::Named(_) => quote! { #where_clause { #(#fields),* } },
108                Fields::Unnamed(_) if is_struct => quote! { (#(#fields),*) #where_clause; },
109                Fields::Unnamed(_) => quote! { (#(#fields),*) },
110                Fields::Unit if is_struct => quote! { ; },
111                Fields::Unit => quote! {},
112            }
113        })
114    };
115    let f_data = |name, this, pre| match &item.data {
116        Data::Struct(_) => {
117            let data = f_data(this, pre, true).next().unwrap();
118            quote_spanned! {item.span()=>
119                #vis struct #name #data
120            }
121        }
122        Data::Enum(data) => {
123            let ident = data.variants.iter().map(|v| &v.ident);
124            let data = f_data(this, pre, false);
125            quote_spanned! {item.span()=>
126                #vis enum #name #where_clause {
127                    #(#ident #data,)*
128                }
129            }
130        }
131        Data::Union(_) => unreachable!(),
132    };
133    let f_def = f_data(
134        quote! { #ident_f<#params> },
135        quote! { #ident<#args> },
136        quote! {},
137    );
138    let fref_def = f_data(
139        quote! { #ident_f_ref<'s_tytro__, #params> },
140        quote! { #ident_ref<'s_tytro__, #args> },
141        quote! { &'s_tytro__ },
142    );
143    let fmut_def = f_data(
144        quote! { #ident_f_mut<'s_tytro__, #params> },
145        quote! { #ident_mut<'s_tytro__, #args> },
146        quote! { &'s_tytro__ mut },
147    );
148
149    let data_enum = variants.iter().map(|(ident, rec, f)| {
150        let f = f
151            .iter()
152            .enumerate()
153            .filter(|(i, _)| Some(*i) != *rec)
154            .map(|(_, f)| &f.ty);
155        if rec.is_some() {
156            quote! { #ident(R_tytro__, #(#f),*) }
157        } else {
158            quote! { #ident(L_tytro__, #(#f),*) }
159        }
160    });
161
162    let variant = |ident: &syn::Ident| {
163        if let Data::Enum(_) = &item.data {
164            quote! { F_tytro__::#ident }
165        } else {
166            quote! { F_tytro__ }
167        }
168    };
169
170    let get_rec_match = variants.iter().map(|(ident, rec, f)| {
171        let Some(rec) = rec else {
172            return quote! {
173                #ident_data::#ident(never, ..) => absurd(never),
174            };
175        };
176        let var = (0..f.len())
177            .filter(|i| i != rec)
178            .map(|i| format_ident!("_{i}"));
179        let f = f.members().enumerate().map(|(i, m)| {
180            if i == *rec {
181                quote! { #m: Ty_tytro__ { last: self.last, rec: more_rec } }
182            } else {
183                let var = format_ident!("_{i}");
184                quote! { #m: #var }
185            }
186        });
187        let variant = variant(ident);
188        quote! {
189            #ident_data::#ident((), #(#var),*) => #variant { #(#f),* },
190        }
191    });
192    let get_rec_match = quote! {
193        match rec {
194            #(#get_rec_match)*
195            #ident_data::__Marker_tytro__(never, ..) => absurd(never),
196        }
197    };
198
199    let get_last_match = variants.iter().map(|(ident, rec, f)| {
200        if rec.is_some() {
201            return quote! {
202                #ident_data::#ident(never, ..) => absurd(never),
203            };
204        }
205        let var = (0..f.len()).map(|i| format_ident!("_{i}"));
206        let var2 = var.clone();
207        let members = f.members();
208        let variant = variant(ident);
209        quote! {
210            #ident_data::#ident((), #(#var),*) => #variant { #(#members: #var2),* },
211        }
212    });
213    let get_last_match = quote! {
214        match self.last {
215            #(#get_last_match)*
216            #ident_data::__Marker_tytro__(never, ..) => absurd(never),
217        }
218    };
219
220    let build_match = variants.iter().map(|(ident, rec, f)| {
221        let variant = variant(ident);
222        if let Some(rec) = *rec {
223            let var = (0..f.len())
224                .filter(|i| *i != rec)
225                .map(|i| format_ident!("_{i}"));
226            let f = f.members().enumerate().map(|(i, m)| {
227                if i == rec {
228                    quote! { #m: rec }
229                } else {
230                    let var = format_ident!("_{i}");
231                    quote! { #m: #var }
232                }
233            });
234            quote! {
235                #variant { #(#f),* } => build_rec(#ident_data::#ident((), #(#var),*), rec),
236            }
237        } else {
238            let var = (0..f.len()).map(|i| format_ident!("_{i}"));
239            let var2 = var.clone();
240            let members = f.members();
241            quote! {
242                #variant { #(#members: #var),* } => build_last(#ident_data::#ident((), #(#var2),*)),
243            }
244        }
245    });
246
247    let rec_ty = quote! {
248        #ident_data<(), ::core::convert::Infallible, #args>
249    };
250    let last_ty = quote! {
251        #ident_data<::core::convert::Infallible, (), #args>
252    };
253
254    let q = quote_spanned! {item.span()=>
255        pub struct #ident<#params> #where_clause {
256            last: #last_ty,
257            rec: ::std::vec::Vec<#rec_ty>,
258        }
259
260        #[derive(Clone, Copy)]
261        pub struct #ident_ref<'s_tytro__, #params> #where_clause {
262            last: &'s_tytro__ #last_ty,
263            rec: &'s_tytro__ [#rec_ty],
264        }
265
266        pub struct #ident_mut<'s_tytro__, #params> #where_clause {
267            last: &'s_tytro__ mut #last_ty,
268            rec: &'s_tytro__ mut [#rec_ty],
269        }
270
271        enum #ident_data<R_tytro__, L_tytro__, #params> {
272            #(#data_enum,)*
273            __Marker_tytro__(::core::convert::Infallible, R_tytro__, L_tytro__),
274        }
275
276        impl<#params_impl> #ident<#args> {
277            pub fn as_ref(&self) -> #ident_ref<'_, #args> {
278                #ident_ref {
279                    last: &self.last,
280                    rec: &self.rec,
281                }
282            }
283
284            pub fn as_mut(&mut self) -> #ident_mut<'_, #args> {
285                #ident_mut {
286                    last: &mut self.last,
287                    rec: &mut self.rec,
288                }
289            }
290
291            pub fn get(self) -> #ident_f<#args> {
292                use #ident as Ty_tytro__;
293                use #ident_f as F_tytro__;
294
295                let mut more_rec = self.rec;
296                let absurd = |never| match never {};
297                match more_rec.pop() {
298                    ::core::option::Option::Some(rec) => #get_rec_match,
299                    _ => #get_last_match,
300                }
301            }
302
303            pub fn get_ref(&self) -> #ident_f_ref<'_, #args> {
304                self.as_ref().get_ref()
305            }
306
307            pub fn get_mut(&mut self) -> #ident_f_mut<'_, #args> {
308                self.as_mut().get_mut()
309            }
310        }
311
312        impl<'s_tytro__, #params_impl> #ident_mut<'s_tytro__, #args> {
313            fn into_ref(self) -> #ident_ref<'s_tytro__, #args> {
314                #ident_ref {
315                    last: self.last,
316                    rec: self.rec,
317                }
318            }
319
320            pub fn as_ref(&self) -> #ident_ref<'_, #args> {
321                #ident_ref {
322                    last: self.last,
323                    rec: self.rec,
324                }
325            }
326
327            pub fn as_mut(&mut self) -> #ident_ref<'_, #args> {
328                #ident_ref {
329                    last: self.last,
330                    rec: self.rec,
331                }
332            }
333
334            pub fn get_ref(self) -> #ident_f_ref<'s_tytro__, #args> {
335                self.into_ref().get_ref()
336            }
337
338            pub fn get_mut(self) -> #ident_f_mut<'s_tytro__, #args> {
339                use #ident_mut as Ty_tytro__;
340                use #ident_f_mut as F_tytro__;
341
342                let absurd = |&mut never| match never {};
343                match self.rec.split_last_mut() {
344                    ::core::option::Option::Some((rec, more_rec)) => #get_rec_match,
345                    _ => #get_last_match,
346                }
347            }
348        }
349
350        impl<'s_tytro__, #params_impl> #ident_ref<'s_tytro__, #args> {
351            pub fn get_ref(self) -> #ident_f_ref<'s_tytro__, #args> {
352                use #ident_ref as Ty_tytro__;
353                use #ident_f_ref as F_tytro__;
354
355                let absurd = |&never| match never {};
356                match self.rec.split_last() {
357                    ::core::option::Option::Some((rec, more_rec)) => #get_rec_match,
358                    _ => #get_last_match,
359                }
360            }
361        }
362
363        impl<#params_impl> #ident_f<#args> {
364            pub fn build(self) -> #ident<#args> {
365                use #ident_f as F_tytro__;
366
367                let build_last = |last| {
368                    let rec = ::std::vec::Vec::new();
369                    #ident { last, rec }
370                };
371                let build_rec = |this, #ident { last, mut rec }| {
372                    rec.push(this);
373                    #ident { last, rec }
374                };
375                match self {
376                    #(#build_match)*
377                }
378            }
379        }
380    };
381    let q = quote_spanned! {item.span()=>
382        #[allow(non_snake_case, non_camel_case_types, unused_variables, dead_code)]
383        mod #ident_mod {
384            use super::*;
385            #q
386        }
387
388        #vis use #ident_mod::{#ident, #ident_ref, #ident_mut};
389
390        #f_def
391        #fref_def
392        #fmut_def
393    };
394    q.into()
395}