Skip to main content

dataloom_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{DeriveInput, Fields, GenericArgument, Ident, parse_macro_input};
4
5use syn::{PathArguments, Type, TypePath};
6
7fn is_option(ty: &Type) -> bool {
8    match ty {
9        Type::Path(TypePath { path, .. }) => {
10            let segment = path.segments.last();
11
12            matches!(segment, Some(seg) if seg.ident == "Option")
13        }
14        _ => false,
15    }
16}
17
18fn option_inner_type(ty: &Type) -> Option<&Type> {
19    let Type::Path(type_path) = ty else {
20        return None;
21    };
22
23    let segment = type_path.path.segments.last()?;
24
25    if segment.ident != "Option" {
26        return None;
27    }
28
29    let PathArguments::AngleBracketed(args) = &segment.arguments else {
30        return None;
31    };
32
33    let GenericArgument::Type(inner_ty) = args.args.first()? else {
34        return None;
35    };
36
37    Some(inner_ty)
38}
39
40fn prefix_ident(ident: Ident) -> Ident {
41    Ident::new(&format!("macro_{ident}"), ident.span())
42}
43
44#[proc_macro_derive(FromIter)]
45pub fn derive_from_iter(input: TokenStream) -> TokenStream {
46    let input = parse_macro_input!(input as DeriveInput);
47
48    if let syn::Data::Struct(ref data) = input.data
49        && let Fields::Named(ref fields) = data.fields
50    {
51        let options = fields.named.iter().map(|field| {
52            let ty = field.ty.clone();
53            let name = field.ident.clone().unwrap();
54            let ty = if is_option(&ty) {
55                option_inner_type(&ty).unwrap().clone()
56            } else {
57                ty
58            };
59
60            let prefixed = prefix_ident(name);
61            quote!(let mut #prefixed: Option<#ty> = None;)
62        });
63
64        let fill_options = fields.named.iter().map(|field| {
65            let name = field.ident.clone().unwrap();
66            let name_string = name.to_string();
67
68            let prefixed = prefix_ident(name);
69            quote!(
70                String { .. } if matches!(Self::get_latest_column_name(#name_string), Some(col) if col == column_name) => {
71                    #prefixed = column_value.from_column(column_type).map_err(|e| DatabaseStrategyError::SearchModel(e.to_string()))?;
72                })
73        });
74
75        let construct_self = fields.named.iter().map(|field| {
76            let name = field.ident.clone().unwrap();
77
78            let prefixed = prefix_ident(name.clone());
79
80            let value = if is_option(&field.ty) {
81                quote!(Some(#prefixed.unwrap()))
82            } else {
83                quote!(#prefixed.unwrap())
84            };
85
86            quote!(
87                #name: #value
88            )
89        });
90
91        let name = input.ident;
92
93        return quote!(
94            impl dataloom::dataloom_db_core::traits::from_iter::FromIter for #name {
95                fn from_iter(iter: impl Iterator<Item = dataloom::dataloom_db_core::traits::from_iter::FromIterValue>) -> Result<Self, dataloom::dataloom_db_core::traits::DatabaseStrategyError>
96                where
97                    Self: Sized,
98                {
99                    use dataloom::dataloom_db_core::{column::{FromColumn, ToColumn}, traits::DatabaseStrategyError};
100                    #(#options)*
101
102                    for dataloom::dataloom_db_core::traits::from_iter::FromIterValue {
103                        column_name,
104                        column_value,
105                        column_type,
106                    } in iter {
107                        match column_name {
108                            #(#fill_options)*
109                            _ => {}
110                        }
111                    }
112
113                    Ok(Self {
114                        #(#construct_self),*
115                    })
116                }
117
118            }
119        )
120        .into();
121    }
122
123    TokenStream::from(
124        syn::Error::new(
125            input.ident.span(),
126            "Only structs with named fields can derive 'FromIter'",
127        )
128        .into_compile_error(),
129    )
130}
131
132#[proc_macro_derive(SaveData)]
133pub fn derive_save_data(input: TokenStream) -> TokenStream {
134    let input = parse_macro_input!(input as DeriveInput);
135
136    if let syn::Data::Struct(ref data) = input.data
137        && let Fields::Named(ref fields) = data.fields
138    {
139        let name = input.ident;
140        // let name_string = name.to_string();
141
142        let save_models = fields.named.iter().map(|field| {
143            let field_name = field.ident.clone().unwrap();
144            let field_name_string = field_name.to_string();
145
146            // let value = if is_option(&field.ty) {
147            //     quote!(self.#field_name)
148            // } else {
149            //     quote!(self.#field_name.clone().into())
150            // };
151
152            quote!(dataloom::dataloom_db_core::save::SaveModel::new(
153                Self::get_latest_column_name(#field_name_string).unwrap(),
154                self.#field_name.to_column().unwrap()
155            ))
156        });
157
158        return quote!(
159            impl dataloom::dataloom_db_core::traits::save_data::SaveData for #name {
160                fn get_save_data(&self) -> Vec<dataloom::dataloom_db_core::save::SaveModel> {
161                    use dataloom::dataloom_db_core::column::ToColumn;
162                    use dataloom::dataloom_db_core::traits::model::Model;
163
164                    vec![
165                        #(#save_models),*
166                    ]
167                }
168
169            }
170        )
171        .into();
172    }
173
174    TokenStream::from(
175        syn::Error::new(
176            input.ident.span(),
177            "Only structs with named fields can derive 'FromIter'",
178        )
179        .into_compile_error(),
180    )
181}