support other collections via #[collection(...)]

1 files changed, 47 insertions, 1 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 13cbce7..05fdb81 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,6 +8,8 @@ use quote::format_ident;
 use quote::quote;
 use quote::quote_spanned;
 use quote::ToTokens;
+use syn::parenthesized;
+use syn::parse::Parse;
 use syn::parse_macro_input;
 use syn::spanned::Spanned;
 use syn::token::Brace;
@@ -16,16 +18,20 @@ use syn::token::Lt;
 use syn::token::Trait;
 use syn::AngleBracketedGenericArguments;
 use syn::Block;
+use syn::Error;
 use syn::Expr;
 use syn::GenericArgument;
 use syn::GenericParam;
+use syn::Ident;
 use syn::ImplItemMethod;
 use syn::ItemTrait;
+use syn::LitInt;
 use syn::Path;
 use syn::PathArguments;
 use syn::PathSegment;
 use syn::Signature;
 use syn::Stmt;
+use syn::Token;
 use syn::TraitBound;
 use syn::TraitItem;
 use syn::TraitItemMethod;
@@ -53,13 +59,37 @@ mod syn_utils;
 mod transform;
 
 macro_rules! abort {
-    ($span:expr, $message:literal $(,$args:tt)*) => {{
+    ($span:expr, $message:literal $(,$args:expr)*) => {{
         let msg = format!($message $(,$args)*);
         let tokens = quote_spanned! {$span => compile_error!(#msg);}.into();
         tokens
     }};
 }
 
+struct Collection {
+    id: Ident,
+    count: usize,
+}
+
+impl Parse for Collection {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let inner;
+        parenthesized!(inner in input);
+
+        let id: Ident = inner.parse()?;
+        let _: Token![,] = inner.parse()?;
+        let count_lit: LitInt = inner.parse()?;
+        let count: usize = count_lit.base10_parse()?;
+        if count < 1 {
+            return Err(Error::new(
+                count_lit.span(),
+                "number of type parameters must be >= 1",
+            ));
+        }
+        Ok(Self { id, count })
+    }
+}
+
 #[proc_macro_attribute]
 pub fn dynamize(_attr: TokenStream, input: TokenStream) -> TokenStream {
     let mut original_trait = parse_macro_input!(input as ItemTrait);
@@ -100,6 +130,22 @@ pub fn dynamize(_attr: TokenStream, input: TokenStream) -> TokenStream {
             };
             let tokens = group.stream();
             dyn_trait_attrs.push(quote! {#[#tokens]});
+        } else if original_trait.attrs[i].path.is_ident("collection") {
+            let attr = original_trait.attrs.remove(i);
+            let tokens = attr.tokens.into();
+            let coll = parse_macro_input!(tokens as Collection);
+
+            if type_converter
+                .collections
+                .insert(coll.id.clone(), coll.count)
+                .is_some()
+            {
+                return abort!(
+                    coll.id.span(),
+                    "collection `{}` is defined multiple times for this trait",
+                    coll.id
+                );
+            }
         } else {
             i += 1;
         }