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).ok();
72                })
73        });
74
75        let check_options = fields.named.iter().map(|field| {
76            let name = field.ident.clone().unwrap();
77            let prefixed = prefix_ident(name.clone());
78            quote!(let Some(#name) = #prefixed)
79        });
80
81        let construct_self = fields.named.iter().map(|field| {
82            let name = field.ident.clone().unwrap();
83
84            let value = if is_option(&field.ty) {
85                quote!(Some(#name))
86            } else {
87                quote!(#name)
88            };
89
90            quote!(
91                #name: #value
92            )
93        });
94
95        let name = input.ident;
96
97        return quote!(
98            impl dataloom::dataloom_db_core::traits::from_iter::FromIter for #name {
99                fn from_iter(iter: impl Iterator<Item = dataloom::dataloom_db_core::traits::from_iter::FromIterValue>) -> Option<Self>
100                where
101                    Self: Sized,
102                {
103                    use dataloom::dataloom_db_core::column::{FromColumn, ToColumn};
104                    #(#options)*
105
106                    for dataloom::dataloom_db_core::traits::from_iter::FromIterValue {
107                        column_name,
108                        column_value,
109                        column_type,
110                    } in iter {
111                        match column_name {
112                            #(#fill_options)*
113                            _ => {}
114                        }
115                    }
116
117                    if #(#check_options)&&* {
118                        return Some(Self {
119                            #(#construct_self),*
120                        });
121                    }
122
123                    None
124                }
125
126            }
127        )
128        .into();
129    }
130
131    TokenStream::from(
132        syn::Error::new(
133            input.ident.span(),
134            "Only structs with named fields can derive 'FromIter'",
135        )
136        .into_compile_error(),
137    )
138}
139
140#[proc_macro_derive(SaveData)]
141pub fn derive_save_data(input: TokenStream) -> TokenStream {
142    let input = parse_macro_input!(input as DeriveInput);
143
144    if let syn::Data::Struct(ref data) = input.data
145        && let Fields::Named(ref fields) = data.fields
146    {
147        let name = input.ident;
148        // let name_string = name.to_string();
149
150        let save_models = fields.named.iter().map(|field| {
151            let field_name = field.ident.clone().unwrap();
152            let field_name_string = field_name.to_string();
153
154            // let value = if is_option(&field.ty) {
155            //     quote!(self.#field_name)
156            // } else {
157            //     quote!(self.#field_name.clone().into())
158            // };
159
160            quote!(dataloom::dataloom_db_core::save::SaveModel::new(
161                Self::get_latest_column_name(#field_name_string).unwrap(),
162                self.#field_name.to_column().unwrap()
163            ))
164        });
165
166        return quote!(
167            impl dataloom::dataloom_db_core::traits::save_data::SaveData for #name {
168                fn get_save_data(&self) -> Vec<dataloom::dataloom_db_core::save::SaveModel> {
169                    use dataloom::dataloom_db_core::column::ToColumn;
170                    use dataloom::dataloom_db_core::traits::model::Model;
171
172                    vec![
173                        #(#save_models),*
174                    ]
175                }
176
177            }
178        )
179        .into();
180    }
181
182    TokenStream::from(
183        syn::Error::new(
184            input.ident.span(),
185            "Only structs with named fields can derive 'FromIter'",
186        )
187        .into_compile_error(),
188    )
189}