1extern crate proc_macro;
32
33mod attr;
34use attr::Attr;
35
36use proc_macro::TokenStream;
37use proc_macro2::Span;
38use proc_macro_error::ResultExt;
39use quote::{quote, ToTokens};
40use syn::{
41 punctuated::Punctuated,
42 Attribute,
43 Data::{Enum, Struct},
44 DataStruct, DeriveInput, Field, Fields, Ident, Token, Variant, Visibility,
45};
46
47#[derive(Clone)]
49struct ItemField<'a> {
50 field: &'a Field,
51 attrs: Vec<Attr>,
52}
53
54impl<'a> ItemField<'a> {
55 fn new(field: &'a Field) -> Self {
56 let attrs = parse_attrs(&field.attrs);
57 Self { field, attrs }
58 }
59
60 fn is_partition_key(&self) -> bool {
61 self.attrs
62 .iter()
63 .any(|attr| matches!(attr, Attr::PartitionKey(_)))
64 }
65
66 fn is_sort_key(&self) -> bool {
67 self.attrs
68 .iter()
69 .any(|attr| matches!(attr, Attr::SortKey(_)))
70 }
71
72 fn is_default_when_absent(&self) -> bool {
73 self.attrs
74 .iter()
75 .any(|attr| matches!(attr, Attr::Default(_)))
76 }
77
78 fn deser_name(&self) -> String {
79 let ItemField { field, attrs } = self;
80 attrs
81 .iter()
82 .find_map(|attr| match attr {
83 Attr::Rename(_, lit) => Some(lit.value()),
84 _ => None,
85 })
86 .unwrap_or_else(|| {
87 field
88 .ident
89 .as_ref()
90 .expect("should have an identifier")
91 .to_string()
92 })
93 }
94}
95
96fn parse_attrs(all_attrs: &[Attribute]) -> Vec<Attr> {
97 all_attrs
98 .iter()
99 .filter(|attr| is_dynomite_attr(attr))
100 .flat_map(|attr| {
101 attr.parse_args_with(Punctuated::<Attr, Token![,]>::parse_terminated)
102 .unwrap_or_abort()
103 })
104 .collect()
105}
106
107#[proc_macro_error::proc_macro_error]
119#[proc_macro_derive(Item, attributes(partition_key, sort_key, dynomite))]
120pub fn derive_item(input: TokenStream) -> TokenStream {
121 let ast = syn::parse_macro_input!(input);
122
123 let gen = match expand_item(ast) {
124 Ok(g) => g,
125 Err(e) => return e.to_compile_error().into(),
126 };
127
128 gen.into_token_stream().into()
129}
130
131#[proc_macro_error::proc_macro_error]
134#[proc_macro_derive(Attributes, attributes(dynomite))]
135pub fn derive_attributes(input: TokenStream) -> TokenStream {
136 let ast = syn::parse_macro_input!(input);
137
138 let gen = match expand_attributes(ast) {
139 Ok(g) => g,
140 Err(e) => return e.to_compile_error().into(),
141 };
142
143 gen.into_token_stream().into()
144}
145
146#[proc_macro_error::proc_macro_error]
152#[proc_macro_derive(Attribute)]
153pub fn derive_attribute(input: TokenStream) -> TokenStream {
154 let ast = syn::parse_macro_input!(input);
155 let gen = expand_attribute(ast);
156 gen.into_token_stream().into()
157}
158
159fn expand_attribute(ast: DeriveInput) -> impl ToTokens {
160 let name = &ast.ident;
161 match ast.data {
162 Enum(variants) => {
163 make_dynomite_attr(name, &variants.variants.into_iter().collect::<Vec<_>>())
164 }
165 _ => panic!("Dynomite Attributes can only be generated for enum types"),
166 }
167}
168
169fn make_dynomite_attr(
190 name: &Ident,
191 variants: &[Variant],
192) -> impl ToTokens {
193 let attr = quote!(::dynomite::Attribute);
194 let err = quote!(::dynomite::AttributeError);
195 let into_match_arms = variants.iter().map(|var| {
196 let vname = &var.ident;
197 quote! {
198 #name::#vname => stringify!(#vname).to_string(),
199 }
200 });
201 let from_match_arms = variants.iter().map(|var| {
202 let vname = &var.ident;
203 quote! {
204 stringify!(#vname) => ::std::result::Result::Ok(#name::#vname),
205 }
206 });
207
208 quote! {
209 impl #attr for #name {
210 fn into_attr(self) -> ::dynomite::dynamodb::AttributeValue {
211 let arm = match self {
212 #(#into_match_arms)*
213 };
214 ::dynomite::dynamodb::AttributeValue {
215 s: ::std::option::Option::Some(arm),
216 ..::std::default::Default::default()
217 }
218 }
219 fn from_attr(value: ::dynomite::dynamodb::AttributeValue) -> ::std::result::Result<Self, #err> {
220 value.s.ok_or(::dynomite::AttributeError::InvalidType)
221 .and_then(|value| match &value[..] {
222 #(#from_match_arms)*
223 _ => ::std::result::Result::Err(::dynomite::AttributeError::InvalidFormat)
224 })
225 }
226 }
227 }
228}
229
230fn expand_attributes(ast: DeriveInput) -> syn::Result<impl ToTokens> {
231 use syn::spanned::Spanned as _;
232 let name = &ast.ident;
233 match ast.data {
234 Struct(DataStruct { fields, .. }) => match fields {
235 Fields::Named(named) => {
236 make_dynomite_attributes(name, &named.named.into_iter().collect::<Vec<_>>())
237 }
238 fields => Err(syn::Error::new(
239 fields.span(),
240 "Dynomite Attributes require named fields",
241 )),
242 },
243 _ => panic!("Dynomite Attributes can only be generated for structs"),
244 }
245}
246
247fn expand_item(ast: DeriveInput) -> syn::Result<impl ToTokens> {
248 use syn::spanned::Spanned as _;
249 let name = &ast.ident;
250 let vis = &ast.vis;
251 match ast.data {
252 Struct(DataStruct { fields, .. }) => match fields {
253 Fields::Named(named) => {
254 make_dynomite_item(vis, name, &named.named.into_iter().collect::<Vec<_>>())
255 }
256 fields => Err(syn::Error::new(
257 fields.span(),
258 "Dynomite Items require named fields",
259 )),
260 },
261 _ => panic!("Dynomite Items can only be generated for structs"),
262 }
263}
264
265fn make_dynomite_attributes(
266 name: &Ident,
267 fields: &[Field],
268) -> syn::Result<impl ToTokens> {
269 let item_fields = fields.iter().map(ItemField::new).collect::<Vec<_>>();
270 let from_attribute_map = get_from_attributes_trait(name, &item_fields)?;
272 let to_attribute_map = get_to_attribute_map_trait(name, &item_fields)?;
274 let attribute = quote!(::dynomite::Attribute);
276 let impl_attribute = quote! {
277 impl #attribute for #name {
278 fn into_attr(self: Self) -> ::dynomite::AttributeValue {
279 ::dynomite::AttributeValue {
280 m: Some(self.into()),
281 ..::dynomite::AttributeValue::default()
282 }
283 }
284 fn from_attr(value: ::dynomite::AttributeValue) -> std::result::Result<Self, ::dynomite::AttributeError> {
285 use ::dynomite::FromAttributes;
286 value
287 .m
288 .ok_or(::dynomite::AttributeError::InvalidType)
289 .and_then(Self::from_attrs)
290 }
291 }
292 };
293
294 Ok(quote! {
295 #from_attribute_map
296 #to_attribute_map
297 #impl_attribute
298 })
299}
300
301fn make_dynomite_item(
302 vis: &Visibility,
303 name: &Ident,
304 fields: &[Field],
305) -> syn::Result<impl ToTokens> {
306 let item_fields = fields.iter().map(ItemField::new).collect::<Vec<_>>();
307 let partition_key_count = item_fields.iter().filter(|f| f.is_partition_key()).count();
309 if partition_key_count != 1 {
310 return Err(syn::Error::new(
311 name.span(),
312 format!(
313 "All Item's must declare one and only one partition_key. The `{}` Item declared {}",
314 name, partition_key_count
315 ),
316 ));
317 }
318 let dynamodb_traits = get_dynomite_item_traits(vis, name, &item_fields)?;
320 let from_attribute_map = get_from_attributes_trait(name, &item_fields)?;
322 let to_attribute_map = get_to_attribute_map_trait(name, &item_fields)?;
324
325 Ok(quote! {
326 #from_attribute_map
327 #to_attribute_map
328 #dynamodb_traits
329 })
330}
331
332fn get_to_attribute_map_trait(
339 name: &Ident,
340 fields: &[ItemField],
341) -> syn::Result<impl ToTokens> {
342 let attributes = quote!(::dynomite::Attributes);
343 let from = quote!(::std::convert::From);
344 let to_attribute_map = get_to_attribute_map_function(name, fields)?;
345
346 Ok(quote! {
347 impl #from<#name> for #attributes {
348 #to_attribute_map
349 }
350 })
351}
352
353fn get_to_attribute_map_function(
365 name: &Ident,
366 fields: &[ItemField],
367) -> syn::Result<impl ToTokens> {
368 let to_attribute_value = quote!(::dynomite::Attribute::into_attr);
369
370 let field_conversions = fields
371 .iter()
372 .map(|field| {
373 let field_deser_name = field.deser_name();
374
375 let field_ident = &field.field.ident;
376 Ok(quote! {
377 values.insert(
378 #field_deser_name.to_string(),
379 #to_attribute_value(item.#field_ident)
380 );
381 })
382 })
383 .collect::<syn::Result<Vec<_>>>()?;
384
385 Ok(quote! {
386 fn from(item: #name) -> Self {
387 let mut values = Self::new();
388 #(#field_conversions)*
389 values
390 }
391 })
392}
393
394fn get_from_attributes_trait(
406 name: &Ident,
407 fields: &[ItemField],
408) -> syn::Result<impl ToTokens> {
409 let from_attrs = quote!(::dynomite::FromAttributes);
410 let from_attribute_map = get_from_attributes_function(fields)?;
411
412 Ok(quote! {
413 impl #from_attrs for #name {
414 #from_attribute_map
415 }
416 })
417}
418
419fn get_from_attributes_function(fields: &[ItemField]) -> syn::Result<impl ToTokens> {
420 let attributes = quote!(::dynomite::Attributes);
421 let from_attribute_value = quote!(::dynomite::Attribute::from_attr);
422 let err = quote!(::dynomite::AttributeError);
423
424 let field_conversions = fields.iter().map(|field| {
425 let field_deser_name = field.deser_name();
427
428 let field_ident = &field.field.ident;
429 if field.is_default_when_absent() {
430 Ok(quote! {
431 #field_ident: match attrs.remove(#field_deser_name) {
432 Some(field) => #from_attribute_value(field)?,
433 _ => ::std::default::Default::default()
434 }
435 })
436 } else {
437 Ok(quote! {
438 #field_ident: #from_attribute_value(
439 attrs.remove(#field_deser_name)
440 .ok_or(::dynomite::AttributeError::MissingField { name: #field_deser_name.to_string() })?
441 )?
442 })
443 }
444 }).collect::<syn::Result<Vec<_>>>()?;
445
446 Ok(quote! {
447 fn from_attrs(mut attrs: #attributes) -> ::std::result::Result<Self, #err> {
448 ::std::result::Result::Ok(Self {
449 #(#field_conversions),*
450 })
451 }
452 })
453}
454
455fn get_dynomite_item_traits(
456 vis: &Visibility,
457 name: &Ident,
458 fields: &[ItemField],
459) -> syn::Result<impl ToTokens> {
460 let impls = get_item_impls(vis, name, fields)?;
461
462 Ok(quote! {
463 #impls
464 })
465}
466
467fn get_item_impls(
468 vis: &Visibility,
469 name: &Ident,
470 fields: &[ItemField],
471) -> syn::Result<impl ToTokens> {
472 let item_trait = get_item_trait(name, fields)?;
474 let key_struct = get_key_struct(vis, name, fields)?;
476
477 Ok(quote! {
478 #item_trait
479 #key_struct
480 })
481}
482
483fn get_item_trait(
493 name: &Ident,
494 fields: &[ItemField],
495) -> syn::Result<impl ToTokens> {
496 let item = quote!(::dynomite::Item);
497 let attribute_map = quote!(
498 ::std::collections::HashMap<String, ::dynomite::dynamodb::AttributeValue>
499 );
500 let partition_key_field = fields.iter().find(|f| f.is_partition_key());
501 let sort_key_field = fields.iter().find(|f| f.is_sort_key());
502 let partition_key_insert = partition_key_field.map(get_key_inserter).transpose()?;
503 let sort_key_insert = sort_key_field.map(get_key_inserter).transpose()?;
504
505 Ok(partition_key_field
506 .map(|_| {
507 quote! {
508 impl #item for #name {
509 fn key(&self) -> #attribute_map {
510 let mut keys = ::std::collections::HashMap::new();
511 #partition_key_insert
512 #sort_key_insert
513 keys
514 }
515 }
516 }
517 })
518 .unwrap_or_else(proc_macro2::TokenStream::new))
519}
520
521fn get_key_inserter(field: &ItemField) -> syn::Result<impl ToTokens> {
527 let to_attribute_value = quote!(::dynomite::Attribute::into_attr);
528
529 let field_deser_name = field.deser_name();
530 let field_ident = &field.field.ident;
531 Ok(quote! {
532 keys.insert(
533 #field_deser_name.to_string(),
534 #to_attribute_value(self.#field_ident.clone())
535 );
536 })
537}
538
539fn get_key_struct(
547 vis: &Visibility,
548 name: &Ident,
549 fields: &[ItemField],
550) -> syn::Result<impl ToTokens> {
551 let name = Ident::new(&format!("{}Key", name), Span::call_site());
552
553 let partition_key_field = fields
554 .iter()
555 .find(|field| field.is_partition_key())
556 .cloned()
557 .map(|field| {
558 let mut field = field.field.clone();
562 field.attrs.retain(is_dynomite_attr);
563
564 quote! {
565 #field
566 }
567 });
568
569 let sort_key_field = fields
570 .iter()
571 .find(|field| field.is_sort_key())
572 .cloned()
573 .map(|field| {
574 let mut field = field.field.clone();
578 field.attrs.retain(is_dynomite_attr);
579
580 quote! {
581 #field
582 }
583 });
584
585 Ok(partition_key_field
586 .map(|partition_key_field| {
587 quote! {
588 #[derive(::dynomite::Attributes, Debug, Clone, PartialEq)]
589 #vis struct #name {
590 #partition_key_field,
591 #sort_key_field
592 }
593 }
594 })
595 .unwrap_or_else(proc_macro2::TokenStream::new))
596}
597
598fn is_dynomite_attr(suspect: &syn::Attribute) -> bool {
599 suspect.path.is_ident("dynomite")
600}