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::{format_ident, 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 -> BigInt },
            diesel_type: quote! { ::tosin::db::sql_types::BigInt },
        }
    } else if field_type == &parse_type("u64") {
        // TODO allow at all since some dbs can't express that type
        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("i64") {
        // TODO default
        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("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 = 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 FieldsInfo {
        tosin_fields,
        diesel_columns,
        diesel_types,
    } = ast_data.fields.iter().map(to_field_spec).collect();
    let non_id_field_names: Vec<_> = ast_data
        .fields
        .iter()
        .filter(|field| !is_id(field))
        .map(|field| &field.ident)
        .collect();

    let model_meta = quote! {
        pub const META: ::tosin::db::models::ModelMeta = ::tosin::db::models::ModelMeta {
            name: stringify!(#name),
            fields: &[ #(#tosin_fields),* ],
        };
    };

    let new = {
        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();
        quote! {
            pub fn new(#(#new_params),*) -> Self {
                Self {
                    id: None,
                    #(#non_id_field_names),*
                }
            }
        }
    };

    let save_mut = quote! {
        pub fn save_mut(&mut self, connection: &tosin::db::backend::Connection) {
            use diesel::prelude::*;
            if self.id.is_none() {
                // no id yet, so not from db, so insert
                // unfortunately InsertStatement::get_result is only supported on pg for now,
                // so we have to pull this shit
                let new_self = tosin::db::backend::insert_and_retrieve::<
                    Self, // row
                    #lowercase_name::table, // table
                    #lowercase_name::id, // id
                >(
                    self,
                    #lowercase_name::table,
                    connection,
                    #lowercase_name::id
                );
                *self = new_self;
            } else {
                diesel::update(self as &Self)
                    .set(self as &Self)
                    .execute(connection)
                    .unwrap();
            }
        }
    };

    let (getters, setters): (Vec<_>, Vec<_>) = ast_data
        .fields
        .iter()
        .map(|field| {
            let ident = field.ident.as_ref().expect("anonymous field in model");
            let ty = &field.ty;
            let getter = quote! {
                pub fn #ident(&self) -> &#ty {
                    &self.#ident
                }
            };
            let set_ident = format_ident!("set_{}", ident);
            // TODO only use impl Into<#ty> sometimes
            let setter = quote! {
                pub fn #set_ident(&mut self, #ident: impl Into<#ty>) {
                    self.#ident = #ident.into();
                }
            };
            (getter, setter)
        })
        .unzip();

    let impl_queryable = {
        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();
        quote! {
            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 (id, #(#non_id_field_names),*): (#(#real_db_types),*) = tosin::db::Queryable::build(row);
                    Self {
                        id: Some(id),
                        #(#non_id_field_names),*
                    }
                }
            }
        }
    };

    let diesel_table = {
        let app_name = {
            // ugly hack because span::module_path() doesn't exist
            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 app_dir = if call_site_path.ends_with("models.rs") {
                call_site_path.parent().unwrap()
            } else {
                todo!("not in models.rs what is this")
            };

            // oh this is so fucking disgusting
            if app_dir.starts_with("examples") {
                app_dir
                    .components()
                    .skip(1) // drop the "examples/", keep the example name
                    .map(|c| c.as_os_str().to_string_lossy())
                    .collect::<Vec<_>>()
                    .join("::")
            } else {
                todo!("what's a {}", call_site_path.display())
            }
        };
        let real_table_name = format!("{}-{}", app_name, lowercase_name);
        quote! {
            // this means users need #[macro_use] extern crate diesel; but fuck doing it ourselves
            table! {
                #[sql_name = #real_table_name]
                #lowercase_name {
                    #(#diesel_columns,)*
                }
            }
        }
    };

    // used for both insertable and aschangeset
    let dsl_eq_types: Vec<_> = ast_data
        .fields
        .iter()
        .filter(|field| !is_id(field))
        .map(|field| {
            let ident = &field.ident;
            let ty = &field.ty;
            quote! {
                Option<tosin::db::dsl::Eq<#lowercase_name::#ident, &'a #ty>>
            }
        })
        .collect();
    let dsl_eq_values: Vec<_> = ast_data
        .fields
        .iter()
        .filter(|field| !is_id(field))
        .map(|field| {
            let ident = &field.ident;
            quote! {
                Some(#lowercase_name::#ident.eq(&self.#ident))
            }
        })
        .collect();

    let impl_insertable = {
        quote! {
            impl<'a> tosin::db::Insertable<#lowercase_name::table> for &'a #name {
                type Values = <(#(#dsl_eq_types,)*) as tosin::db::Insertable<#lowercase_name::table>>::Values;

                fn values(self) -> Self::Values {
                    use diesel::ExpressionMethods;
                    (#(#dsl_eq_values,)*).values()
                }
            }
        }
    };

    let impl_identifiable = {
        quote! {
            impl<'a> tosin::db::Identifiable for &'a #name {
                type Id = Id;
                fn id(self) -> Id {
                    // TODO this probably isn't good
                    self.id.expect("not actually identifiable if you haven't saved yet")
                }
            }
        }
    };

    let impl_has_table = {
        quote! {
            impl tosin::db::associations::HasTable for #name {
                type Table = #lowercase_name::table;
                fn table() -> #lowercase_name::table {
                    #lowercase_name::table
                }
            }
        }
    };

    let impl_as_changeset = {
        quote! {
            impl<'a> tosin::db::query_builder::AsChangeset for &'a #name {
                type Target = #lowercase_name::table;
                type Changeset = <(#(#dsl_eq_types,)*) as tosin::db::query_builder::AsChangeset>::Changeset;

                fn as_changeset(self) -> Self::Changeset {
                    use diesel::ExpressionMethods;
                    (#(#dsl_eq_values,)*).as_changeset()
                }
            }
        }
    };

    let gen = quote! {
        impl #name {
            #model_meta

            #new

            #save_mut

            #(#getters)*
            #(#setters)*
        }

        #diesel_table
        #impl_queryable
        #impl_insertable
        #impl_identifiable
        #impl_has_table
        #impl_as_changeset
    };
    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()
}