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

2 files changed, 52 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;
         }
diff --git a/src/transform.rs b/src/transform.rs
index 05866b6..11a98c2 100644
--- a/src/transform.rs
+++ b/src/transform.rs
@@ -17,6 +17,7 @@ use crate::{
 #[derive(Default)]
 pub struct TypeConverter<'a> {
     pub assoc_type_conversions: HashMap<Ident, DestType<'a>>,
+    pub collections: HashMap<Ident, usize>,
 }
 
 #[derive(Debug)]
@@ -39,6 +40,10 @@ impl TypeConverter<'_> {
     /// ... etc. A return type of None means the type isn't recognized.
     #[rustfmt::skip]
     fn get_collection_type_count(&self, ident: &Ident) -> Option<usize> {
+        if let Some(count) = self.collections.get(ident) {
+            return Some(*count);
+        }
+
         // when adding a type here don't forget to document it in the README
         if ident == "Vec"        { return Some(1); }
         if ident == "VecDeque"   { return Some(1); }