iridis_message_derive/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{
6    parse_macro_input, punctuated::Punctuated, token::Comma, DeriveInput, Field, Fields, Ident,
7    Token, Variant,
8};
9
10#[proc_macro_derive(ArrowMessage)]
11pub fn from_into_arrow_derive(input: TokenStream) -> TokenStream {
12    let input = parse_macro_input!(input as DeriveInput);
13
14    let name = input.ident;
15
16    match input.data {
17        syn::Data::Struct(s) => match s.fields {
18            Fields::Named(ref fields) => struct_derive(name, fields.named.clone()),
19            _ => panic!("Only structs with named fields are supported"),
20        },
21        syn::Data::Enum(e) => enum_derive(name, e.variants.clone()),
22        _ => panic!("Only structs and enums are supported"),
23    }
24}
25
26fn struct_derive(name: Ident, fields: Punctuated<Field, Comma>) -> TokenStream {
27    let field_attributes = fields
28        .iter()
29        .map(|field| (&field.ident, &field.ty))
30        .collect::<Vec<_>>();
31
32    let fields_fill = field_attributes.iter().map(|&(field, ty)| {
33        quote! {
34            <#ty>::field(stringify!(#field)),
35        }
36    });
37
38    let union_data_fill = field_attributes.iter().map(|&(field, _)| {
39        quote! {
40            #field: extract_union_data(stringify!(#field), &map, &children)?,
41        }
42    });
43
44    let arrow_data_fill = field_attributes.iter().map(|&(field, _)| {
45        quote! {
46            self.#field.try_into_arrow()?,
47        }
48    });
49
50    let expanded = quote! {
51        impl ArrowMessage for #name {
52            fn field(name: impl Into<String>) -> iridis_message::prelude::thirdparty::arrow_schema::Field {
53                make_union_fields(
54                    name,
55                    vec![
56                        #(#fields_fill)*
57                    ],
58                )
59            }
60
61            fn try_from_arrow(data: iridis_message::prelude::thirdparty::arrow_data::ArrayData) -> iridis_message::prelude::thirdparty::eyre::Result<Self>
62            where
63                Self: Sized,
64            {
65                let (map, children) = unpack_union(data);
66
67                Ok(Self {
68                    #(#union_data_fill)*
69                })
70            }
71
72            fn try_into_arrow(self) -> iridis_message::prelude::thirdparty::eyre::Result<iridis_message::prelude::thirdparty::arrow_array::ArrayRef> {
73                let union_fields = get_union_fields::<Self>()?;
74
75                make_union_array(
76                    union_fields,
77                    vec![
78                        #(#arrow_data_fill)*
79                    ],
80                )
81            }
82        }
83
84        impl TryFrom<iridis_message::prelude::thirdparty::arrow_data::ArrayData> for #name {
85            type Error = iridis_message::prelude::thirdparty::eyre::Report;
86
87            fn try_from(data: iridis_message::prelude::thirdparty::arrow_data::ArrayData) -> iridis_message::prelude::thirdparty::eyre::Result<Self> {
88                #name::try_from_arrow(data)
89            }
90        }
91
92        impl TryFrom<#name> for iridis_message::prelude::thirdparty::arrow_data::ArrayData {
93            type Error = iridis_message::prelude::thirdparty::eyre::Report;
94
95            fn try_from(item: #name) -> iridis_message::prelude::thirdparty::eyre::Result<Self> {
96                use iridis_message::prelude::thirdparty::arrow_array::Array;
97
98                item.try_into_arrow().map(|array| array.into_data())
99            }
100        }
101    };
102
103    TokenStream::from(expanded)
104}
105
106fn enum_derive(name: Ident, variants: Punctuated<Variant, Token![,]>) -> TokenStream {
107    let variants: Vec<_> = variants
108        .iter()
109        .map(|variant| {
110            let variant_name = &variant.ident;
111            let variant_str = variant_name.to_string().to_lowercase(); // Exemple : `Foo` -> "foo"
112            (variant_name, variant_str)
113        })
114        .collect();
115
116    let into_string_arms = variants.iter().map(|(variant_name, variant_str)| {
117        quote! {
118            #name::#variant_name => #variant_str.to_string(),
119        }
120    });
121
122    let try_from_string_arms = variants.iter().map(|(variant_name, variant_str)| {
123        quote! {
124            #variant_str => Ok(#name::#variant_name),
125        }
126    });
127
128    let expanded = quote! {
129        impl #name {
130            pub fn into_string(self) -> String {
131                match self {
132                    #(#into_string_arms)*
133                }
134            }
135
136            pub fn try_from_string(s: String) -> iridis_message::prelude::thirdparty::eyre::Result<Self> {
137                match s.as_str() {
138                    #(#try_from_string_arms)*
139                    _ => Err(iridis_message::prelude::thirdparty::eyre::eyre!("Invalid value for {}: {}", stringify!(#name), s)),
140                }
141            }
142        }
143
144        impl ArrowMessage for #name {
145            fn field(name: impl Into<String>) -> iridis_message::prelude::thirdparty::arrow_schema::Field {
146                String::field(name)
147            }
148
149            fn try_from_arrow(data: iridis_message::prelude::thirdparty::arrow_data::ArrayData) -> iridis_message::prelude::thirdparty::eyre::Result<Self>
150            where
151                Self: Sized,
152            {
153                Encoding::try_from_string(String::try_from_arrow(data)?)
154            }
155
156            fn try_into_arrow(self) -> iridis_message::prelude::thirdparty::eyre::Result<iridis_message::prelude::thirdparty::arrow_array::ArrayRef> {
157                String::try_into_arrow(self.into_string())
158            }
159        }
160
161
162        impl TryFrom<iridis_message::prelude::thirdparty::arrow_data::ArrayData> for #name {
163            type Error = iridis_message::prelude::thirdparty::eyre::Report;
164
165            fn try_from(data: iridis_message::prelude::thirdparty::arrow_data::ArrayData) -> iridis_message::prelude::thirdparty::eyre::Result<Self> {
166                #name::try_from_arrow(data)
167            }
168        }
169
170        impl TryFrom<#name> for iridis_message::prelude::thirdparty::arrow_data::ArrayData {
171            type Error = iridis_message::prelude::thirdparty::eyre::Report;
172
173            fn try_from(item: #name) -> iridis_message::prelude::thirdparty::eyre::Result<Self> {
174                item.try_into_arrow().map(|array| array.into_data())
175            }
176        }
177    };
178
179    TokenStream::from(expanded)
180}
iridis_message_derive/lib.rs

iridis_message_derive/
lib.rs
