path: root/tosin-macros/src/lib.rs

#![feature(proc_macro_span)]

extern crate proc_macro;

use std::collections::HashMap;
use std::fs;

use proc_macro::TokenStream;
use quote::quote;

#[proc_macro_derive(Model, attributes(model))]
pub fn model_derive(input: TokenStream) -> TokenStream {
    // Construct a representation of Rust code as a syntax tree
    // that we can manipulate
    let ast = syn::parse(input).unwrap();

    // Build the trait implementation
    impl_model(&ast)
}

struct FieldInfo<T: quote::ToTokens> {
    tosin_field: T,
    diesel_column: T,
    diesel_type: T,
}

struct FieldsInfo<T: quote::ToTokens> {
    tosin_fields: Vec<T>,
    diesel_columns: Vec<T>,
    diesel_types: Vec<T>,
}

impl<T: quote::ToTokens> std::iter::FromIterator<FieldInfo<T>> for FieldsInfo<T> {
    fn from_iter<I: IntoIterator<Item = FieldInfo<T>>>(iter: I) -> Self {
        let mut result = Self {
            tosin_fields: vec![],
            diesel_columns: vec![],
            diesel_types: vec![],
        };
        for info in iter {
            result.tosin_fields.push(info.tosin_field);
            result.diesel_columns.push(info.diesel_column);
            result.diesel_types.push(info.diesel_type);
        }
        result
    }
}

fn to_field_spec(field: &syn::Field) -> FieldInfo<impl quote::ToTokens> {
    fn parse_type(ty: &str) -> syn::Type {
        syn::parse_str(ty).unwrap()
    }
    let field_name = &field.ident;
    let field_type = &field.ty;
    let model_options: HashMap<syn::Path, syn::Lit> = field
        .attrs
        .iter()
        .filter_map(|attr| attr.parse_meta().ok())
        .filter_map(|meta| {
            if let syn::Meta::List(meta) = meta {
                Some(meta)
            } else {
                None
            }
        })
        .filter(|meta| meta.path.get_ident().map_or(false, |path| path == "model"))
        .flat_map(|model| {
            model.nested.into_iter().filter_map(|item| {
                if let syn::NestedMeta::Meta(syn::Meta::NameValue(data)) = item {
                    Some((data.path, data.lit))
                } else {
                    None
                }
            })
        })
        .collect();
    if field_type == &parse_type("Option<Id>") {
        FieldInfo {
            tosin_field: quote! { ::tosin::db::models::Field::IntField { name: stringify!(#field_name) } },
            diesel_column: quote! { #field_name -> BigInt },
            diesel_type: quote! { ::tosin::db::sql_types::BigInt },
        }
    } else if field_type == &parse_type("Id") {
        // TODO foreign key constraint
        FieldInfo {
            tosin_field: quote! { ::tosin::db::models::Field::IntField { name: stringify!(#field_name) } },
            diesel_column: quote! { #field_name -> Integer },
            diesel_type: quote! { ::tosin::db::sql_types::Integer },
        }
    } else if field_type == &parse_type("usize") {
        // TODO default
        FieldInfo {
            tosin_field: quote! { ::tosin::db::models::Field::IntField { name: stringify!(#field_name) } },
            diesel_column: quote! { #field_name -> Integer },
            diesel_type: quote! { ::tosin::db::sql_types::Integer },
        }
    } else if field_type == &parse_type("String") {
        let max_length = model_options
            .iter()
            .find(|(name, _value)| name.get_ident().map_or(false, |path| path == "max_length"))
            .map(|(_name, value)| value);
        if let Some(max_length) = max_length {
            FieldInfo {
                tosin_field: quote! { ::tosin::db::models::Field::CharField { name: stringify!(#field_name), max_length: Some(#max_length) } },
                diesel_column: quote! { #field_name -> Text },
                diesel_type: quote! { ::tosin::db::sql_types::Text },
            }
        } else {
            FieldInfo {
                tosin_field: quote! { ::tosin::db::models::Field::CharField { name: stringify!(#field_name), max_length: None } },
                diesel_column: quote! { #field_name -> Text },
                diesel_type: quote! { ::tosin::db::sql_types::Text },
            }
        }
    } else if field_type == &parse_type("chrono::NaiveDateTime") {
        FieldInfo {
            tosin_field: quote! { ::tosin::db::models::Field::DateTimeField { name: stringify!(#field_name) } },
            diesel_column: quote! { #field_name -> Timestamp },
            diesel_type: quote! { ::tosin::db::sql_types::Timestamp },
        }
    } else {
        use quote::ToTokens;
        panic!("can't handle {}", field.to_token_stream())
    }
}

fn is_id(field: &syn::Field) -> bool {
    let name_matches = field
        .ident
        .as_ref()
        .map_or(false, |name| name.to_string() == "id");
    let type_matches = field.ty == syn::parse_str("Option<Id>").unwrap();
    name_matches && type_matches
}

fn impl_model(ast: &syn::DeriveInput) -> TokenStream {
    let name = &ast.ident;
    let lowercase_name = quote::format_ident!("{}", name.to_string().to_lowercase());
    let ast_data = if let syn::Data::Struct(ast_data) = &ast.data {
        ast_data
    } else {
        panic!("not on a struct");
    };
    let real_db_types: Vec<_> = ast_data
        .fields
        .iter()
        .map(|field| {
            if is_id(field) {
                syn::parse_str("Id").unwrap()
            } else {
                field.ty.clone()
            }
        })
        .collect();
    let new_params: Vec<_> = ast_data
        .fields
        .iter()
        .filter(|field| !is_id(field))
        .map(|field| {
            let ty = &field.ty;
            let name = &field.ident;
            quote! { #name: #ty }
        })
        .collect();
    let field_names: Vec<_> = ast_data
        .fields
        .iter()
        .filter(|field| !is_id(field))
        .map(|field| &field.ident)
        .collect();
    let FieldsInfo {
        tosin_fields,
        diesel_columns,
        diesel_types,
    } = ast_data.fields.iter().map(to_field_spec).collect();
    let gen = quote! {
        impl #name {
            pub const META: ::tosin::db::models::ModelMeta = ::tosin::db::models::ModelMeta {
                name: stringify!(#name),
                fields: &[ #(#tosin_fields),* ],
            };

            pub fn new(#(#new_params),*) -> Self {
                Self {
                    id: None,
                    #(#field_names),*
                }
            }
        }

        impl<__DB: tosin::db::diesel_backend::Backend, __ST> tosin::db::Queryable<__ST, __DB> for #name
        where (#(#real_db_types),*): tosin::db::Queryable<__ST, __DB> {
            type Row = <(#(#real_db_types),*) as tosin::db::Queryable<__ST, __DB>>::Row;

            fn build(row: Self::Row) -> Self {
                let row: (#(#real_db_types),*) = tosin::db::Queryable::build(row);
                todo!()
            }
        }

        // this means users need #[macro_use] extern crate diesel; but fuck doing it ourselves
        table! {
            #lowercase_name {
                #(#diesel_columns,)*
            }
        }
    };
    gen.into()
}

#[proc_macro]
pub fn gather_migrations(_input: TokenStream) -> TokenStream {
    let call_site = proc_macro::Span::call_site();
    let call_site_file = call_site.source_file();
    let call_site_path = call_site_file.path();
    if !call_site_file.is_real() {
        panic!("call site does not have a real path");
    }

    let migrations_dir = call_site_path.parent().unwrap();
    let migrations: Vec<syn::Ident> = migrations_dir
        .read_dir()
        .unwrap()
        .map(Result::unwrap)
        .map(|x| x.path().file_stem().unwrap().to_string_lossy().into_owned())
        .filter(|x| x != "mod")
        .map(|x| syn::parse_str(&x).unwrap())
        .collect();

    let gen = quote! {
        #( mod #migrations; )*

        pub const ALL: &[Migration] = &[
            #(#migrations::MIGRATION),*
        ];
    };

    gen.into()
}

#[proc_macro]
pub fn gather_models(_input: TokenStream) -> TokenStream {
    let call_site = proc_macro::Span::call_site();
    let call_site_file = call_site.source_file();
    let call_site_path = call_site_file.path();
    if !call_site_file.is_real() {
        panic!("call site does not have a real path");
    }

    let call_site_ast = syn::parse_file(&fs::read_to_string(call_site_path).unwrap()).unwrap();
    let models = call_site_ast
        .items
        .iter()
        .filter_map(|item| {
            if let syn::Item::Struct(item) = item {
                Some(item)
            } else {
                None
            }
        })
        .filter(|item| {
            item.attrs.iter().any(|attr| {
                let attr = if let Ok(syn::Meta::List(attr)) = attr.parse_meta() {
                    attr
                } else {
                    return false;
                };
                if attr
                    .path
                    .get_ident()
                    .map_or(false, |hopefully_derive| hopefully_derive == "derive")
                {
                    let mut derived = attr.nested.iter().filter_map(|derived| {
                        if let syn::NestedMeta::Meta(derived) = derived {
                            Some(derived)
                        } else {
                            None
                        }
                    });
                    derived.any(|derived| {
                        derived
                            .path()
                            .get_ident()
                            .map_or(false, |hopefully_model| hopefully_model == "Model")
                    })
                } else {
                    false
                }
            })
        })
        .map(|item| &item.ident);

    let gen = quote! {
        pub const ALL: &[tosin::db::models::ModelMeta] = &[
            #(#models::META),*
        ];
    };

    gen.into()
}