tytro 0.1.0

类型上的尾递归优化
Documentation
#![doc = include_str!("../README.md")]

use proc_macro::TokenStream;
use quote::{ToTokens, format_ident, quote, quote_spanned};
use syn::{Data, Field, Fields, GenericParam, parse_macro_input, spanned::Spanned};

#[proc_macro_attribute]
pub fn tytro(_cfg: TokenStream, item: TokenStream) -> TokenStream {
    let item = parse_macro_input!(item as syn::DeriveInput);

    assert!(item.attrs.is_empty(), "attribute is unsupported");

    let find_rec = |fields: &Fields| {
        let mut found = None;
        for (i, f) in fields.iter().enumerate() {
            assert!(f.attrs.is_empty(), "attribute is unsupported");
            match &f.ty {
                syn::Type::Path(p) if p.path.is_ident("Self") => {}
                _ => continue,
            }
            let old = found.replace(i);
            assert!(
                old.is_none(),
                "Self can only appear at most once per variant"
            );
        }
        found
    };

    let variants = match &item.data {
        Data::Struct(data) => vec![(&item.ident, find_rec(&data.fields), &data.fields)],
        Data::Enum(data) => Vec::from_iter(data.variants.iter().map(|v| {
            assert!(v.discriminant.is_none(), "discriminant is unsupported");
            assert!(v.attrs.is_empty(), "attribute is unsupported");
            (&v.ident, find_rec(&v.fields), &v.fields)
        })),
        Data::Union(_) => panic!("union is unsupported"),
    };

    assert!(
        variants.iter().any(|(_, rec, _)| rec.is_some()),
        "Self must appear at least once"
    );

    let syn::DeriveInput {
        vis,
        ident,
        generics:
            syn::Generics {
                params,
                where_clause,
                ..
            },
        ..
    } = &item;

    let ident_mod = format_ident!("__{}__mod_tytro__", ident);
    let ident_ref = format_ident!("{}Ref", ident);
    let ident_mut = format_ident!("{}Mut", ident);
    let ident_f = format_ident!("{}F", ident);
    let ident_f_ref = format_ident!("{}FRef", ident);
    let ident_f_mut = format_ident!("{}FMut", ident);
    let ident_data = format_ident!("__{}__data_tytro__", ident);

    let where_clause = where_clause.iter().flat_map(|c| c.predicates.iter());
    let where_clause = quote! {
        where #(#where_clause,)*
    };

    let args = params.iter().map::<&dyn ToTokens, _>(|param| match param {
        GenericParam::Lifetime(p) => &p.lifetime,
        GenericParam::Type(p) => &p.ident,
        GenericParam::Const(p) => &p.ident,
    });
    let args = quote! { #(#args),* };

    let params_impl = params.iter().map(|param| {
        let mut param = param.clone();
        match &mut param {
            GenericParam::Lifetime(_) => {}
            GenericParam::Type(param) => param.default = None,
            GenericParam::Const(param) => param.default = None,
        }
        param
    });
    let params_impl = quote! { #(#params_impl),* };

    let f_data = |this, pre, is_struct| {
        let where_clause = &where_clause;
        variants.iter().map(move |(_, rec, f)| {
            let fields = f.iter().enumerate();
            let fields = fields.map(|(i, f)| {
                let Field {
                    vis,
                    ident,
                    colon_token,
                    ty,
                    ..
                } = f;
                if Some(i) == *rec {
                    quote! { #vis #ident #colon_token #this }
                } else {
                    quote! { #vis #ident #colon_token #pre #ty }
                }
            });
            match f {
                Fields::Named(_) => quote! { #where_clause { #(#fields),* } },
                Fields::Unnamed(_) if is_struct => quote! { (#(#fields),*) #where_clause; },
                Fields::Unnamed(_) => quote! { (#(#fields),*) },
                Fields::Unit if is_struct => quote! { ; },
                Fields::Unit => quote! {},
            }
        })
    };
    let f_data = |name, this, pre| match &item.data {
        Data::Struct(_) => {
            let data = f_data(this, pre, true).next().unwrap();
            quote_spanned! {item.span()=>
                #vis struct #name #data
            }
        }
        Data::Enum(data) => {
            let ident = data.variants.iter().map(|v| &v.ident);
            let data = f_data(this, pre, false);
            quote_spanned! {item.span()=>
                #vis enum #name #where_clause {
                    #(#ident #data,)*
                }
            }
        }
        Data::Union(_) => unreachable!(),
    };
    let f_def = f_data(
        quote! { #ident_f<#params> },
        quote! { #ident<#args> },
        quote! {},
    );
    let fref_def = f_data(
        quote! { #ident_f_ref<'s_tytro__, #params> },
        quote! { #ident_ref<'s_tytro__, #args> },
        quote! { &'s_tytro__ },
    );
    let fmut_def = f_data(
        quote! { #ident_f_mut<'s_tytro__, #params> },
        quote! { #ident_mut<'s_tytro__, #args> },
        quote! { &'s_tytro__ mut },
    );

    let data_enum = variants.iter().map(|(ident, rec, f)| {
        let f = f
            .iter()
            .enumerate()
            .filter(|(i, _)| Some(*i) != *rec)
            .map(|(_, f)| &f.ty);
        if rec.is_some() {
            quote! { #ident(R_tytro__, #(#f),*) }
        } else {
            quote! { #ident(L_tytro__, #(#f),*) }
        }
    });

    let variant = |ident: &syn::Ident| {
        if let Data::Enum(_) = &item.data {
            quote! { F_tytro__::#ident }
        } else {
            quote! { F_tytro__ }
        }
    };

    let get_rec_match = variants.iter().map(|(ident, rec, f)| {
        let Some(rec) = rec else {
            return quote! {
                #ident_data::#ident(never, ..) => absurd(never),
            };
        };
        let var = (0..f.len())
            .filter(|i| i != rec)
            .map(|i| format_ident!("_{i}"));
        let f = f.members().enumerate().map(|(i, m)| {
            if i == *rec {
                quote! { #m: Ty_tytro__ { last: self.last, rec: more_rec } }
            } else {
                let var = format_ident!("_{i}");
                quote! { #m: #var }
            }
        });
        let variant = variant(ident);
        quote! {
            #ident_data::#ident((), #(#var),*) => #variant { #(#f),* },
        }
    });
    let get_rec_match = quote! {
        match rec {
            #(#get_rec_match)*
            #ident_data::__Marker_tytro__(never, ..) => absurd(never),
        }
    };

    let get_last_match = variants.iter().map(|(ident, rec, f)| {
        if rec.is_some() {
            return quote! {
                #ident_data::#ident(never, ..) => absurd(never),
            };
        }
        let var = (0..f.len()).map(|i| format_ident!("_{i}"));
        let var2 = var.clone();
        let members = f.members();
        let variant = variant(ident);
        quote! {
            #ident_data::#ident((), #(#var),*) => #variant { #(#members: #var2),* },
        }
    });
    let get_last_match = quote! {
        match self.last {
            #(#get_last_match)*
            #ident_data::__Marker_tytro__(never, ..) => absurd(never),
        }
    };

    let build_match = variants.iter().map(|(ident, rec, f)| {
        let variant = variant(ident);
        if let Some(rec) = *rec {
            let var = (0..f.len())
                .filter(|i| *i != rec)
                .map(|i| format_ident!("_{i}"));
            let f = f.members().enumerate().map(|(i, m)| {
                if i == rec {
                    quote! { #m: rec }
                } else {
                    let var = format_ident!("_{i}");
                    quote! { #m: #var }
                }
            });
            quote! {
                #variant { #(#f),* } => build_rec(#ident_data::#ident((), #(#var),*), rec),
            }
        } else {
            let var = (0..f.len()).map(|i| format_ident!("_{i}"));
            let var2 = var.clone();
            let members = f.members();
            quote! {
                #variant { #(#members: #var),* } => build_last(#ident_data::#ident((), #(#var2),*)),
            }
        }
    });

    let rec_ty = quote! {
        #ident_data<(), ::core::convert::Infallible, #args>
    };
    let last_ty = quote! {
        #ident_data<::core::convert::Infallible, (), #args>
    };

    let q = quote_spanned! {item.span()=>
        pub struct #ident<#params> #where_clause {
            last: #last_ty,
            rec: ::std::vec::Vec<#rec_ty>,
        }

        #[derive(Clone, Copy)]
        pub struct #ident_ref<'s_tytro__, #params> #where_clause {
            last: &'s_tytro__ #last_ty,
            rec: &'s_tytro__ [#rec_ty],
        }

        pub struct #ident_mut<'s_tytro__, #params> #where_clause {
            last: &'s_tytro__ mut #last_ty,
            rec: &'s_tytro__ mut [#rec_ty],
        }

        enum #ident_data<R_tytro__, L_tytro__, #params> {
            #(#data_enum,)*
            __Marker_tytro__(::core::convert::Infallible, R_tytro__, L_tytro__),
        }

        impl<#params_impl> #ident<#args> {
            pub fn as_ref(&self) -> #ident_ref<'_, #args> {
                #ident_ref {
                    last: &self.last,
                    rec: &self.rec,
                }
            }

            pub fn as_mut(&mut self) -> #ident_mut<'_, #args> {
                #ident_mut {
                    last: &mut self.last,
                    rec: &mut self.rec,
                }
            }

            pub fn get(self) -> #ident_f<#args> {
                use #ident as Ty_tytro__;
                use #ident_f as F_tytro__;

                let mut more_rec = self.rec;
                let absurd = |never| match never {};
                match more_rec.pop() {
                    ::core::option::Option::Some(rec) => #get_rec_match,
                    _ => #get_last_match,
                }
            }

            pub fn get_ref(&self) -> #ident_f_ref<'_, #args> {
                self.as_ref().get_ref()
            }

            pub fn get_mut(&mut self) -> #ident_f_mut<'_, #args> {
                self.as_mut().get_mut()
            }
        }

        impl<'s_tytro__, #params_impl> #ident_mut<'s_tytro__, #args> {
            fn into_ref(self) -> #ident_ref<'s_tytro__, #args> {
                #ident_ref {
                    last: self.last,
                    rec: self.rec,
                }
            }

            pub fn as_ref(&self) -> #ident_ref<'_, #args> {
                #ident_ref {
                    last: self.last,
                    rec: self.rec,
                }
            }

            pub fn as_mut(&mut self) -> #ident_ref<'_, #args> {
                #ident_ref {
                    last: self.last,
                    rec: self.rec,
                }
            }

            pub fn get_ref(self) -> #ident_f_ref<'s_tytro__, #args> {
                self.into_ref().get_ref()
            }

            pub fn get_mut(self) -> #ident_f_mut<'s_tytro__, #args> {
                use #ident_mut as Ty_tytro__;
                use #ident_f_mut as F_tytro__;

                let absurd = |&mut never| match never {};
                match self.rec.split_last_mut() {
                    ::core::option::Option::Some((rec, more_rec)) => #get_rec_match,
                    _ => #get_last_match,
                }
            }
        }

        impl<'s_tytro__, #params_impl> #ident_ref<'s_tytro__, #args> {
            pub fn get_ref(self) -> #ident_f_ref<'s_tytro__, #args> {
                use #ident_ref as Ty_tytro__;
                use #ident_f_ref as F_tytro__;

                let absurd = |&never| match never {};
                match self.rec.split_last() {
                    ::core::option::Option::Some((rec, more_rec)) => #get_rec_match,
                    _ => #get_last_match,
                }
            }
        }

        impl<#params_impl> #ident_f<#args> {
            pub fn build(self) -> #ident<#args> {
                use #ident_f as F_tytro__;

                let build_last = |last| {
                    let rec = ::std::vec::Vec::new();
                    #ident { last, rec }
                };
                let build_rec = |this, #ident { last, mut rec }| {
                    rec.push(this);
                    #ident { last, rec }
                };
                match self {
                    #(#build_match)*
                }
            }
        }
    };
    let q = quote_spanned! {item.span()=>
        #[allow(non_snake_case, non_camel_case_types, unused_variables, dead_code)]
        mod #ident_mod {
            use super::*;
            #q
        }

        #vis use #ident_mod::{#ident, #ident_ref, #ident_mut};

        #f_def
        #fref_def
        #fmut_def
    };
    q.into()
}