derive_dynamic_node/
lib.rs

1#![recursion_limit = "128"]
2extern crate proc_macro;
3
4use proc_macro::TokenStream;
5use proc_macro2::TokenStream as TS2;
6use quote::{quote, ToTokens};
7use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, Type};
8
9/// #[derive(DynamicNode)]
10/// pub struct Ec2Instance2 {
11///     arn: String,
12///     launch_time: u64
13/// }
14
15fn name_and_ty(field: &Field) -> (&Ident, &Type) {
16    (field.ident.as_ref().unwrap(), &field.ty)
17}
18
19#[proc_macro_derive(DynamicNode)]
20pub fn derive_dynamic_node(input: TokenStream) -> TokenStream {
21    let input: syn::DeriveInput = syn::parse_macro_input!(input as syn::DeriveInput);
22
23    let input_struct = match input.data {
24        Data::Struct(input_struct) => input_struct,
25        _ => panic!("Only available for struct"),
26    };
27
28    let fields = match input_struct.fields {
29        Fields::Named(fields) => fields.named,
30        _ => panic!("Requires named fields"),
31    };
32
33    let methods = fields
34        .iter()
35        .map(name_and_ty)
36        .map(|(name, ty)| get_method(name, ty))
37        .fold(quote!(), |mut acc, method| {
38            acc.extend(method);
39            acc
40        });
41
42    let struct_name = &input.ident;
43    let struct_name_string = input.ident.to_string();
44
45    let node_name = format!("{}Node", struct_name);
46    let node_name = syn::Ident::new(&node_name, struct_name.span());
47
48    let node_trait_name = format!("I{}Node", struct_name);
49    let node_trait_name = syn::Ident::new(&node_trait_name, struct_name.span());
50
51    let q = quote!(
52
53        #[derive(Clone, Debug)]
54        pub struct #node_name {
55            dynamic_node: grapl_graph_descriptions::graph_description::DynamicNode,
56        }
57
58        pub trait #node_trait_name {
59            fn get_mut_dynamic_node(&mut self) -> &mut DynamicNode;
60
61            #methods
62        }
63
64        impl #node_name {
65            pub fn new(strategy: grapl_graph_descriptions::graph_description::IdStrategy, seen_at: u64) -> Self {
66                let mut properties = std::collections::HashMap::with_capacity(1);
67
68                let dynamic_node = DynamicNode {
69                    asset_id: None,
70                    host_ip: None,
71                    hostname: None,
72                    node_type: #struct_name_string .to_owned(),
73                    id_strategy: vec![strategy],
74                    node_key: uuid::Uuid::new_v4().to_string(),
75                    properties,
76                    seen_at,
77                };
78
79                Self { dynamic_node }
80            }
81
82            pub fn with_asset_id(&mut self, asset_id: impl Into<Option<String>>) -> &mut Self {
83                let asset_id = asset_id.into();
84                self.dynamic_node.asset_id = asset_id.clone();
85                if let Some(asset_id) = asset_id {
86                    self.dynamic_node.properties.insert("asset_id".to_owned(), asset_id.into());
87                }
88                self
89            }
90
91            pub fn get_node_key(&self) -> &str {
92                &self.dynamic_node.node_key
93            }
94
95            pub fn clone_node_key(&self) -> String {
96                self.dynamic_node.node_key.clone()
97            }
98
99            pub fn into_dyn_node(self) -> DynamicNode {
100                self.dynamic_node
101            }
102        }
103
104        impl AsRef<grapl_graph_descriptions::graph_description::DynamicNode> for #node_name {
105            fn as_ref(&self) -> &DynamicNode {
106                &self.dynamic_node
107            }
108        }
109
110        impl AsMut<grapl_graph_descriptions::graph_description::DynamicNode> for #node_name {
111            fn as_mut(&mut self) -> &mut DynamicNode {
112                &mut self.dynamic_node
113            }
114        }
115
116        impl Into<grapl_graph_descriptions::graph_description::DynamicNode> for #node_name {
117            fn into(self) -> DynamicNode {
118                self.dynamic_node
119            }
120        }
121
122        impl Into<grapl_graph_descriptions::graph_description::Node> for #node_name {
123            fn into(self) -> Node {
124                self.dynamic_node.into()
125            }
126        }
127
128
129        impl Into<grapl_graph_descriptions::graph_description::Node> for & #node_name {
130            fn into(self) -> Node {
131                self.dynamic_node.clone().into()
132            }
133        }
134
135
136        impl Into<grapl_graph_descriptions::graph_description::Node> for &mut #node_name {
137            fn into(self) -> Node {
138                self.dynamic_node.clone().into()
139            }
140        }
141
142    );
143
144    q.into()
145
146}
147
148
149#[proc_macro_derive(GraplStaticId, attributes(grapl))]
150pub fn derive_grapl_session(input: TokenStream) -> TokenStream {
151    let input: syn::DeriveInput = syn::parse_macro_input!(input as syn::DeriveInput);
152
153    let input_struct = match input.data {
154        Data::Struct(input_struct) => input_struct,
155        _ => panic!("Only available for struct"),
156    };
157
158    let fields = match input_struct.fields {
159        Fields::Named(fields) => fields.named,
160        _ => panic!("Requires named fields"),
161    };
162
163    let id_fields = fields
164        .iter()
165        .filter_map(|field| {
166            for attr in &field.attrs {
167                if attr.path.segments.is_empty() {
168                    return None
169                }
170
171                let id = &attr.path.segments[0].ident;
172                if id.to_string() != "grapl" {
173                    continue
174                }
175
176                return field.ident.as_ref()
177            }
178
179            None
180        })
181        .fold(quote!(), |mut acc, f| {
182            let f = f.to_string();
183            acc.extend(quote!(#f .to_string(), ));
184            acc
185        });
186    
187    assert!(id_fields.to_string().len() > 0);
188
189    let struct_name = &input.ident;
190
191    let node_name = format!("{}Node", struct_name);
192    let node_name = syn::Ident::new(&node_name, struct_name.span());
193
194
195    let strategy_name = format!("{}strategy", struct_name);
196    let strategy_name = syn::Ident::new(&strategy_name, struct_name.span());
197
198    let q = quote!(
199
200        impl #node_name {
201            pub fn static_strategy() -> IdStrategy {
202                Static {
203                    primary_key_properties: vec![
204                        #id_fields
205                    ],
206                    primary_key_requires_asset_id: false,
207                }.into()
208            }
209        }
210    );
211
212
213    q.into()
214
215}
216
217
218fn get_method(property_name: &Ident, property_type: &Type) -> TS2 {
219    let method_name = format!("with_{}", property_name);
220    let method_name = syn::Ident::new(&method_name, property_name.span());
221
222    let property_name_str = format!("{}", property_name);
223    quote!(
224            fn #method_name(&mut self, #property_name: impl Into<#property_type>) -> &mut Self {
225                let #property_name = #property_name .into();
226                self.get_mut_dynamic_node()
227                .properties.insert(
228                    #property_name_str .to_string(),
229                    #property_name .into(),
230                );
231                self
232            }
233        )
234}
235
236#[cfg(test)]
237mod tests {
238    #[test]
239    fn it_works() {
240        assert_eq!(2 + 2, 4);
241    }
242}