rust-rel8-derive 0.2.1

use darling::{FromDeriveInput, FromField, FromTypeParam, ast, util::Flag};
use proc_macro2::TokenStream;
use quote::quote;
use syn::{DeriveInput, parse_macro_input};

#[derive(FromTypeParam, Debug)]
#[darling(attributes(table))]
struct GenericOpts {
    ident: syn::Ident,
    bounds: Vec<syn::TypeParamBound>,

    /// Use this to mark the type param as being kept in the derived instances.
    ///
    /// These should come after the `'scope` lifetime and the `Mode` type param.
    proxy: Flag,
}

#[derive(FromField, Debug)]
#[darling(attributes(table))]
struct FieldOpts {
    ident: Option<syn::Ident>,

    /// Marks this field as a nested table, and not a column.
    ///
    /// Use this when a field is another table struct.
    nested: Flag,
}

#[derive(FromDeriveInput, Debug)]
#[darling(attributes(table), supports(struct_named))]
struct TableStructOpts {
    ident: syn::Ident,
    data: ast::Data<darling::util::Ignored, FieldOpts>,
    generics: darling::ast::Generics<darling::ast::GenericParam<GenericOpts>>,

    #[darling(rename = "crate", default)]
    crate_name: Option<syn::Path>,
}

/// Derive the necessary traits for Table structs
///
/// This must be used on a struct with named fields (for now).
/// This must be used on a struct which has type parameters in the form: `<'scope, Mode: TableMode, ...>`
///
/// This derive macro has the following attributes:
///
/// - `#[table(proxy)]`: Used on any type parameter coming after `Mode`, this
///     tells us to wire them up. (This will be improved in the future)
///
/// - `#[table(nested)]`: Used on a field which has the type `OtherTable<'scope, Mode>`,
///     instructs the derive macro to treat this field as nested.
#[proc_macro_derive(TableStruct, attributes(table))]
pub fn derive_table_struct(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let a = match TableStructOpts::from_derive_input(&parse_macro_input!(input as DeriveInput)) {
        Ok(a) => a,
        Err(e) => {
            return proc_macro::TokenStream::from(e.write_errors());
        }
    };

    let r = match do_all(&a) {
        Ok(r) => r,
        Err(e) => {
            return proc_macro::TokenStream::from(e.write_errors());
        }
    };

    r.into()
}

fn do_all(a: &TableStructOpts) -> darling::Result<TokenStream> {
    #[cfg(feature = "sqlx")]
    let sqlx_table_loader = do_sqlx_table_loader(a)?;
    #[cfg(not(feature = "sqlx"))]
    let sqlx_table_loader = quote!();

    let rest = do_rest(a)?;

    Ok(quote! {
        #sqlx_table_loader
        #rest
    })
}

#[cfg(feature = "sqlx")]
fn do_sqlx_table_loader(t: &TableStructOpts) -> darling::Result<TokenStream> {
    use proc_macro2::Span;

    let crate_ = t.crate_name.clone().unwrap_or_else(|| {
        let mut path = syn::Path::from(syn::Ident::new("rust_rel8", Span::call_site()));
        path.leading_colon = Some(syn::Token![::](Span::call_site()));
        path
    });
    let ident = &t.ident;
    let proxied_type_params = t
        .generics
        .type_params()
        .filter(|p| p.proxy.is_present())
        .map(|p| p.ident.clone())
        .collect::<Vec<_>>();
    let proxied_type_param_bounds = t
        .generics
        .type_params()
        .filter(|p| p.proxy.is_present())
        .flat_map(|p| {
            p.bounds.iter().map(|b| {
                let ident = &p.ident;
                quote! { #ident: #b}
            })
        })
        .collect::<Vec<_>>();
    let tokens = quote! {
        impl<#(#proxied_type_params,)*> #crate_::TableLoaderSqlx for #ident<'static, #crate_::table_modes::ExprMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            fn load<'a>(
                &self,
                values: &mut impl Iterator<Item = ::sqlx::postgres::PgValueRef<'a>>,
            ) -> Self::Result {
                #crate_::TableUsingMapper::wrap_ref(self).load(values)
            }

            fn skip<'a>(&self, values: &mut impl Iterator<Item = ::sqlx::postgres::PgValueRef<'a>>) {
                #crate_::TableUsingMapper::wrap_ref(self).skip(values)
            }
        }

        impl<#(#proxied_type_params,)*> #crate_::TableLoaderManySqlx for #ident<'static, #crate_::table_modes::ExprMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            fn load_many<'a>(
                &self,
                values: &mut impl Iterator<Item = ::sqlx::postgres::PgValueRef<'a>>,
            ) -> Vec<Self::Result> {
                #crate_::TableUsingMapper::wrap_ref(self).load_many(values)
            }
        }
    };

    Ok(tokens)
}

fn do_rest(t: &TableStructOpts) -> darling::Result<TokenStream> {
    use proc_macro2::Span;

    let crate_ = t.crate_name.clone().unwrap_or_else(|| {
        let mut path = syn::Path::from(syn::Ident::new("rust_rel8", Span::call_site()));
        path.leading_colon = Some(syn::Token![::](Span::call_site()));
        path
    });
    let ident = &t.ident;
    let proxied_type_params = t
        .generics
        .type_params()
        .filter(|p| p.proxy.is_present())
        .map(|p| p.ident.clone())
        .collect::<Vec<_>>();
    let proxied_type_param_bounds = t
        .generics
        .type_params()
        .filter(|p| p.proxy.is_present())
        .flat_map(|p| {
            p.bounds.iter().map(|b| {
                let ident = &p.ident;
                quote! { #ident: #b}
            })
        })
        .collect::<Vec<_>>();

    let fields = t.data.as_ref().take_struct().unwrap();

    let fields = fields
        .iter()
        .map(|f| {
            let ident = f.ident.as_ref().unwrap();
            let nested = f.nested.is_present();

            (ident, nested)
        })
        .collect::<Vec<_>>();

    let map_modes_final_fields = fields.iter().map(|(ident, _)| quote! { #ident });
    let map_modes_final = quote! {
        #ident {
            #(#map_modes_final_fields,)*
        }
    };

    let map_modes_fields = fields.iter().map(|(ident, is_nested)| {
        if *is_nested {
            quote! { let #ident = self.#ident.map_modes(mapper); }
        } else {
            quote! { let #ident = mapper.map_mode(self.#ident); }
        }
    });

    let map_modes_ref_fields = fields.iter().map(|(ident, is_nested)| {
        if *is_nested {
            quote! { let #ident = self.#ident.map_modes_ref(mapper); }
        } else {
            quote! { let #ident = mapper.map_mode_ref(&self.#ident); }
        }
    });

    let map_modes_mut_fields = fields.iter().map(|(ident, is_nested)| {
        if *is_nested {
            quote! { let #ident = self.#ident.map_modes_mut(mapper); }
        } else {
            quote! { let #ident = mapper.map_mode_mut(&mut self.#ident); }
        }
    });

    let with_lt_fields_with_lt = fields
        .iter()
        .map(|(ident, _is_nested)| {
            quote! { #ident: self.#ident.with_lt(marker) }
        })
        .collect::<Vec<_>>();

    let unwith_lt_fields_with_lt = fields
        .iter()
        .map(|(ident, _is_nested)| {
            quote! { #ident: #crate_::ForLifetimeTable::unwith_lt(with_lt.#ident, marker) }
        })
        .collect::<Vec<_>>();

    let shorten_lifetime_fields_shorten = fields
        .iter()
        .map(|(ident, _is_nested)| {
            quote! { #ident: self.#ident.shorten_lifetime() }
        })
        .collect::<Vec<_>>();

    let shorten_lifetime_fields_noop = fields
        .iter()
        .map(|(ident, is_nested)| {
            if *is_nested {
                quote! { #ident: self.#ident.shorten_lifetime() }
            } else {
                quote! { #ident: self.#ident }
            }
        })
        .collect::<Vec<_>>();

    let tokens = quote! {
        impl<'scope, #(#proxied_type_params,)*> #crate_::ForLifetimeTable for #ident<'scope, #crate_::ExprMode, #(#proxied_type_params,)*>
        where
            for<'a> #ident<'a, #crate_::ExprMode, #(#proxied_type_params,)*>: #crate_::Table,
            #(#proxied_type_param_bounds,)*
        {
            type WithLt<'lt> = #ident<'lt, #crate_::ExprMode, #(#proxied_type_params,)*>;

            fn with_lt<'lt>(self, marker: &mut #crate_::WithLtMarker) -> Self::WithLt<'lt> {
                #ident {
                    #(#with_lt_fields_with_lt,)*
                }
            }

            fn unwith_lt<'lt>(with_lt: Self::WithLt<'lt>, marker: &mut #crate_::WithLtMarker) -> Self {
                #ident {
                    #(#unwith_lt_fields_with_lt,)*
                }
            }
        }

        impl<'scope, #(#proxied_type_params,)*> #crate_::ShortenLifetime for #ident<'scope, #crate_::NameMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            type Shortened<'small> = #ident<'small, #crate_::NameMode, #(#proxied_type_params,)*>
            where
                Self: 'small
                ;

            fn shorten_lifetime<'small, 'large: 'small>(self) -> Self::Shortened<'small>
            where
                Self: 'large,
            {
                #ident {
                    #(#shorten_lifetime_fields_noop,)*
                }
            }
        }

        impl<'scope, #(#proxied_type_params,)*> #crate_::ShortenLifetime for #ident<'scope, #crate_::ValueMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            type Shortened<'small> = #ident<'small, #crate_::ValueMode, #(#proxied_type_params,)*>
            where
                Self: 'small
                ;

            fn shorten_lifetime<'small, 'large: 'small>(self) -> Self::Shortened<'small>
            where
                Self: 'large,
            {
                #ident {
                    #(#shorten_lifetime_fields_noop,)*
                }
            }
        }

        impl<'scope, #(#proxied_type_params,)*> #crate_::ShortenLifetime for #ident<'scope, #crate_::EmptyMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            type Shortened<'small> = #ident<'small, #crate_::EmptyMode, #(#proxied_type_params,)*>
            where
                Self: 'small
                ;

            fn shorten_lifetime<'small, 'large: 'small>(self) -> Self::Shortened<'small>
            where
                Self: 'large,
            {
                #ident {
                    #(#shorten_lifetime_fields_noop,)*
                }
            }
        }

        impl<'scope, #(#proxied_type_params,)*> #crate_::ShortenLifetime for #ident<'scope, #crate_::ExprMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            type Shortened<'small> = #ident<'small, #crate_::ExprMode, #(#proxied_type_params,)*>
            where
                Self: 'small
                ;

            fn shorten_lifetime<'small, 'large: 'small>(self) -> Self::Shortened<'small>
            where
                Self: 'large,
            {
                #ident {
                    #(#shorten_lifetime_fields_shorten,)*
                }
            }
        }

        impl<'scope, T: #crate_::TableMode, #(#proxied_type_params,)*> #crate_::TableHKT for #ident<'scope, T, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            type InMode<Mode: #crate_::TableMode> = #ident<'scope, Mode, #(#proxied_type_params,)*>;

            type Mode = T;
        }

        impl<'scope, Mode: #crate_::TableMode, #(#proxied_type_params,)*> #crate_::MapTable<'scope> for #ident<'scope, Mode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            fn map_modes<Mapper, DestMode>(self, mapper: &mut Mapper) -> Self::InMode<DestMode>
            where
                Mapper: #crate_::ModeMapper<'scope, Self::Mode, DestMode>,
                DestMode: #crate_::TableMode,
            {
                #(#map_modes_fields)*

                #map_modes_final
            }

            fn map_modes_ref<Mapper, DestMode>(&self, mapper: &mut Mapper) -> Self::InMode<DestMode>
            where
                Mapper: #crate_::ModeMapperRef<'scope, Self::Mode, DestMode>,
                DestMode: #crate_::TableMode,
            {
                #(#map_modes_ref_fields)*

                #map_modes_final
            }

            fn map_modes_mut<Mapper, DestMode>(&mut self, mapper: &mut Mapper) -> Self::InMode<DestMode>
            where
                Mapper: #crate_::ModeMapperMut<'scope, Self::Mode, DestMode>,
                DestMode: #crate_::TableMode,
            {
                #(#map_modes_mut_fields)*

                #map_modes_final
            }
        }

        impl<'scope, #(#proxied_type_params,)*> #crate_::Table for #ident<'scope, #crate_::table_modes::ExprMode, #(#proxied_type_params,)*>
        where
            #(#proxied_type_param_bounds,)*
        {
            type Result = <Self as #crate_::TableHKT>::InMode<#crate_::table_modes::ValueMode>;

            fn visit(&self, f: &mut impl FnMut(&#crate_::ErasedExpr), mode: VisitTableMode) {
                #crate_::TableUsingMapper::wrap_ref(self).visit(f, mode)
            }

            fn visit_mut(&mut self, f: &mut impl FnMut(&mut #crate_::ErasedExpr), mode: VisitTableMode) {
                #crate_::TableUsingMapper::wrap_mut(self).visit_mut(f, mode)
            }
        }
    };

    Ok(tokens)
}