microrm/microrm-macros/src/lib.rs

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, DeriveInput};

use convert_case::{Case, Casing};

fn parse_microrm_ref(attrs: &[syn::Attribute]) -> proc_macro2::TokenStream {
    for attr in attrs {
        if attr.path.segments.is_empty() {
            continue;
        }

        if attr.tokens.is_empty() && attr.path.segments.last().unwrap().ident == "microrm_internal"
        {
            return quote! { crate };
        }
    }

    quote! { ::microrm }
}

fn parse_fk(attrs: &[syn::Attribute]) -> bool {
    for attr in attrs {
        if attr.path.segments.len() == 1
            && attr.path.segments.last().unwrap().ident == "microrm_foreign"
        {
            return true;
        }
    }

    false
}

/// Turns a serializable/deserializable struct into a microrm entity model.
///
/// There are two important visible effects:
/// - Provides an implementation of `microrm::model::Entity`
/// - Defines a <struct-name>Columns enum
///
/// Note that names are converted from CamelCase to snake_case and vice versa
/// where applicable, so a struct named `TestModel` is given a table name `test_model`
/// and a struct field named `field_name` is given a variant name of `FieldName`.
///
/// The `#[microrm...]` attributes can be used to control the derivation somewhat.
/// The following are understood for the Entity struct:
/// - `#[microrm_internal]`: this is internal to the microrm crate (of extremely limited usefulness
/// outside the microrm library)
///
/// The following are understood on individual fields:
/// - `#[microrm_foreign]`: this is a foreign key (and the field must be of a type implementing `EntityID`)
#[proc_macro_derive(Entity, attributes(microrm_internal, microrm_foreign))]
pub fn derive_entity(tokens: TokenStream) -> TokenStream {
    let input = parse_macro_input!(tokens as DeriveInput);

    let microrm_ref = parse_microrm_ref(&input.attrs);

    let struct_name = &input.ident;
    let enum_name = format_ident!("{}Columns", &input.ident);
    let id_name = format_ident!("{}ID", &input.ident);

    let table_name = format!("{}", struct_name).to_case(Case::Snake);

    let st = match input.data {
        syn::Data::Struct(st) => st,
        _ => panic!("Can only use derive(Entity) on structs!"),
    };
    let fields = match st.fields {
        syn::Fields::Named(fields) => fields,
        _ => panic!("Can only use derive(Entity) on non-unit structs with named fields!"),
    };

    let mut variants = Vec::new();
    let mut field_names = Vec::new();
    let mut field_numbers = Vec::new();
    let mut field_types = Vec::new();
    let mut value_references = Vec::new();

    let mut foreign_keys = Vec::new();
    let mut foreign_key_impls = Vec::new();

    for name in fields.named.iter() {
        let converted_case =
            format!("{}", name.ident.as_ref().unwrap().clone()).to_case(Case::UpperCamel);
        let converted_case = format_ident!("{}", converted_case);
        variants.push(converted_case.clone());

        let field_name = name.ident.as_ref().unwrap().clone();
        let field_name_str = format!("{}", field_name);
        field_names.push(quote! { Self::Column::#converted_case => #field_name_str });

        let nn = field_numbers.len() + 1;
        field_numbers.push(quote! { #nn => Self::#converted_case, });

        let ty = &name.ty;
        field_types.push(quote!{ <#ty as #microrm_ref::model::Modelable>::column_type() });

        if parse_fk(&name.attrs) {
            let fk_struct_name = format_ident!("{}{}ForeignKey", struct_name, converted_case);
            let ty = &name.ty;
            foreign_keys.push(quote! {
                &#fk_struct_name { col: #enum_name::#converted_case }
            });
            foreign_key_impls.push(quote!{
                struct #fk_struct_name {
                    col: #enum_name
                }
                impl #microrm_ref::model::EntityForeignKey<#enum_name> for #fk_struct_name {
                    fn local_column(&self) -> &#enum_name { &self.col }
                    fn foreign_table_name(&self) -> &'static str {
                        <<#ty as #microrm_ref::model::EntityID>::Entity as #microrm_ref::model::Entity>::table_name()
                    }
                    fn foreign_column_name(&self) -> &'static str {
                        "id"
                    }
                }
            });
        }

        value_references.push(quote! { &self. #field_name });
    }

    let column_types_name = format_ident!("{}_COLUMN_TYPES", struct_name.to_string().to_case(Case::ScreamingSnake));

    let field_count = fields.named.iter().count();

    quote!{
        // Related types for #struct_name
        #[derive(Clone,Copy,PartialEq,Hash)]
        #[allow(unused)]
        #[repr(usize)]
        pub enum #enum_name {
            ID,
            #(#variants),*
        }

        #[derive(Debug,PartialEq,Clone,Copy,#microrm_ref::re_export::serde::Serialize,#microrm_ref::re_export::serde::Deserialize)]
        #[allow(unused)]
        pub struct #id_name (i64);

        // Implementations for related types
        impl #microrm_ref::model::EntityColumns for #enum_name {
            type Entity = #struct_name;
        }

        impl std::convert::From<usize> for #enum_name {
            fn from(i: usize) -> Self {
                match i {
                    0 => Self::ID,
                    #(#field_numbers)*
                    _ => {
                        panic!("Given invalid usize to convert to column")
                    },
                }
            }
        }

        impl #microrm_ref::model::EntityID for #id_name {
            type Entity = #struct_name;
            fn from_raw_id(raw: i64) -> Self { Self(raw) }
            fn raw_id(&self) -> i64 { self.0 }
        }

        impl #microrm_ref::model::Modelable for #id_name {
            fn bind_to(&self, stmt: &mut #microrm_ref::re_export::sqlite::Statement, col: usize) -> #microrm_ref::re_export::sqlite::Result<()> {
                use #microrm_ref::re_export::sqlite::Bindable;
                self.0.bind(stmt, col)
            }
            fn build_from(stmt: &#microrm_ref::re_export::sqlite::Statement, col_offset: usize) -> #microrm_ref::re_export::sqlite::Result<(Self, usize)> where Self: Sized {
                stmt.read::<i64>(col_offset).map(|x| (#id_name(x), 1))
            }
            fn column_type() -> &'static str where Self: Sized {
                "integer"
            }
        }
        #microrm_ref::re_export::lazy_static::lazy_static!{
            static ref #column_types_name: [&'static str;#field_count + 1] = {
                [
                    "id",
                    #(#field_types),*
                ]
            };
        }

        // Implementations for #struct_name
        impl #microrm_ref::model::Entity for #struct_name {
            type Column = #enum_name;
            type ID = #id_name;

            fn table_name() -> &'static str { #table_name }
            fn column_count() -> usize {
                // +1 for ID column
                #field_count + 1
            }
            fn index(c: Self::Column) -> usize {
                c as usize
            }
            fn name(c: Self::Column) -> &'static str {
                match c {
                    Self::Column::ID => "ID",
                    #(#field_names),*
                }
            }
            fn values(&self) -> Vec<&dyn #microrm_ref::model::Modelable> {
                vec![ #(#value_references),* ]
            }
            fn column_types() -> &'static [&'static str] {
                #column_types_name.as_ref()
            }
            fn foreign_keys() -> &'static [&'static dyn #microrm_ref::model::EntityForeignKey<Self::Column>] {
                &[#(#foreign_keys),*]
            }
        }

        // Foreign key struct implementations
        #(#foreign_key_impls)*
    }.into()
}

/// Marks a struct or enum as able to be directly used in an Entity to correspond to a single database column.
#[proc_macro_derive(Modelable, attributes(microrm_internal))]
pub fn derive_modelable(tokens: TokenStream) -> TokenStream {
    let input = parse_macro_input!(tokens as DeriveInput);

    if let syn::Data::Enum(e) = &input.data {
        return derive_modelable_enum(&input, e)
    }

    if let syn::Data::Struct(s) = &input.data {
        if s.fields.len() == 1 {
            return derive_transparent_struct(&input, s)
        }
    }

    derive_modelable_general(&input)
}

fn derive_transparent_struct(input: &syn::DeriveInput, ds: &syn::DataStruct) -> TokenStream {
    // for single-element structs, we can simply store these transparently however the element
    // would be stored
    let microrm_ref = parse_microrm_ref(&input.attrs);
    let struct_name = &input.ident;

    let field = ds.fields.iter().next().unwrap();

    let (bind_to, build_from) =
        if let Some(i) = &field.ident {
            (
                quote!{ self.#i.bind_to(stmt, col) },
                quote!{ Self { #i: field.0 } }
            )
        }
        else {
            (
                quote!{ self.0.bind_to(stmt, col) },
                quote!{ Self(field.0) }
            )
        };

    let field_ty = &field.ty;

    quote! {
        impl #microrm_ref::model::Modelable for #struct_name {
            fn bind_to(&self, stmt: &mut #microrm_ref::re_export::sqlite::Statement, col: usize) -> #microrm_ref::re_export::sqlite::Result<()> {
                #bind_to
            }

            fn build_from(stmt: &#microrm_ref::re_export::sqlite::Statement, col_offset: usize) -> #microrm_ref::re_export::sqlite::Result<(Self,usize)> {
                let field = #field_ty::build_from(stmt, col_offset)?;
                Ok((#build_from, field.1))
            }

            fn column_type() -> &'static str where Self: Sized {
                <#field_ty as #microrm_ref::model::Modelable>::column_type()
            }
        }
    }.into()
}

fn derive_modelable_enum(input: &syn::DeriveInput, de: &syn::DataEnum) -> TokenStream {
    for variant in &de.variants {
        if !variant.fields.is_empty() {
            return derive_modelable_general(input)
        }
    }

    // only unit variants! we can store as a string
    let microrm_ref = parse_microrm_ref(&input.attrs);
    let enum_name = &input.ident;

    let mut variant_names = Vec::new();
    let mut variant_name_strs = Vec::new();

    for variant in &de.variants {
        variant_names.push(variant.ident.clone());
        variant_name_strs.push(variant.ident.to_string());
    }

    quote! {
        impl #microrm_ref::model::Modelable for #enum_name {
            fn bind_to(&self, stmt: &mut #microrm_ref::re_export::sqlite::Statement, col: usize) -> #microrm_ref::re_export::sqlite::Result<()> {
                match self {
                    #(
                        Self::#variant_names => #variant_name_strs.bind_to(stmt, col)
                    ),*
                }
            }

            fn build_from(stmt: &#microrm_ref::re_export::sqlite::Statement, col_offset: usize) -> #microrm_ref::re_export::sqlite::Result<(Self,usize)> {
                let str_form = String::build_from(stmt, col_offset)?.0;
                #(
                    if str_form == #variant_name_strs { return Ok((Self::#variant_names, 1)) }
                )*

                return Err(#microrm_ref::re_export::sqlite::Error { code: None, message: None })
            }

            fn column_type() -> &'static str where Self: Sized {
                "text"
            }
        }
    }.into()
}

fn derive_modelable_general(input: &syn::DeriveInput) -> TokenStream {
    let microrm_ref = parse_microrm_ref(&input.attrs);

    let ident = &input.ident;

    quote!{
        impl #microrm_ref::model::Modelable for #ident {
            fn bind_to(&self, stmt: &mut #microrm_ref::re_export::sqlite::Statement, col: usize) -> #microrm_ref::re_export::sqlite::Result<()> {
                use #microrm_ref::re_export::sqlite;
                use #microrm_ref::model::Modelable;
                serde_json::to_string(self).expect("can be serialized").bind_to(stmt, col)
            }
            fn build_from(stmt: &#microrm_ref::re_export::sqlite::Statement, col_offset: usize) -> #microrm_ref::re_export::sqlite::Result<(Self,usize)> {
                use #microrm_ref::re_export::sqlite;
                use #microrm_ref::model::Modelable;
                let str_data = stmt.read::<String>(col_offset).map_err(|e| sqlite::Error { code: None, message: Some(e.to_string()) })?;
                let data = serde_json::from_str(str_data.as_str()).map_err(|e| sqlite::Error { code: None, message: Some(e.to_string()) })?;
                Ok((data,1))
            }
            fn column_type() -> &'static str where Self: Sized {
                "blob"
            }
        }
    }.into()
}

type ColumnList = syn::punctuated::Punctuated<syn::TypePath, syn::Token![,]>;
struct MakeIndexParams {
    unique: Option<syn::Token![!]>,
    name: syn::Ident,
    #[allow(dead_code)]
    comma: syn::Token![,],
    columns: ColumnList,
}

impl syn::parse::Parse for MakeIndexParams {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        Ok(Self {
            unique: input.parse()?,
            name: input.parse()?,
            comma: input.parse()?,
            columns: ColumnList::parse_separated_nonempty(input)?,
        })
    }
}

fn do_make_index(tokens: TokenStream, microrm_ref: proc_macro2::TokenStream) -> TokenStream {
    let input = parse_macro_input!(tokens as MakeIndexParams);

    let index_struct_name = input.name;

    let first_col = input.columns.first().unwrap();
    let mut column_type_path = first_col.path.clone();

    // remove variant name
    column_type_path.segments.pop();
    let last = column_type_path
        .segments
        .pop()
        .expect("Full path to EntityColumn variant");
    column_type_path.segments.push(last.value().clone());

    let index_entity_type_name = format_ident!("{}Entity", index_struct_name);
    let columns = input.columns.clone().into_iter();

    let index_sql_name = format!("{}", index_struct_name).to_case(Case::Snake);

    let unique = input.unique.is_some();

    quote!{
        pub struct #index_struct_name {}
        type #index_entity_type_name = <#column_type_path as #microrm_ref::model::EntityColumns>::Entity;

        impl #microrm_ref::model::Index for #index_struct_name {
            type IndexedEntity = #index_entity_type_name;
            fn index_name() -> &'static str {
                #index_sql_name
            }
            fn columns() -> &'static [#column_type_path] where Self: Sized {
                &[#(#columns),*]
            }
            fn unique() -> bool where Self: Sized {
                #unique
            }
        }
    }.into()
}

/// Defines a struct to represent a optionally-unique index on a table.
///
/// Suppose the following `Entity` definition is used:
///
/// ```ignore
/// #[derive(Entity,Serialize,Deserialize)]
/// struct SystemUser {
///     username: String,
///     hashed_password: String
/// }
/// ```
///
/// We can now use `make_index!` to define an index on the username field:
/// ```ignore
/// make_index!(SystemUsernameIndex, SystemUserColumns::Username)
/// ```
///
/// This index can be made unique by adding a `!` prior to the type name, as:
/// ```ignore
/// make_index!(!SystemUsernameUniqueIndex, SystemUserColumns::Username)
/// ```
#[proc_macro]
pub fn make_index(tokens: TokenStream) -> TokenStream {
    do_make_index(tokens, quote! { ::microrm })
}

/// For internal use inside the microrm library. See `make_index`.
#[proc_macro]
pub fn make_index_internal(tokens: TokenStream) -> TokenStream {
    do_make_index(tokens, quote! { crate })
}

// , attributes(microrm_internal))]