generalize to collection with one type / two types

3 files changed, 32 insertions, 29 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 8fa1294..3b4aeeb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -411,14 +411,14 @@ impl TypeTransform {
                     quote! {#arg.map_err(|x| #err_inner)}
                 }
             }
-            TypeTransform::Collection(inner) => {
-                let inner = inner.convert(quote!(x));
-                quote! {#arg.into_iter().map(|x| #inner).collect()}
+            TypeTransform::CollectionOneType(type1) => {
+                let type1 = type1.convert(quote!(v1));
+                quote! {#arg.into_iter().map(|v1| #type1).collect()}
             }
-            TypeTransform::CollectionMap(key, value) => {
-                let key = key.convert(quote!(k));
-                let value = value.convert(quote!(v));
-                quote! {#arg.into_iter().map(|(k,v)| (#key, #value)).collect()}
+            TypeTransform::CollectionTwoTypes(type1, type2) => {
+                let type1 = type1.convert(quote!(v1));
+                let type2 = type2.convert(quote!(v2));
+                quote! {#arg.into_iter().map(|(v1,v2)| (#type1, #type2)).collect()}
             }
             TypeTransform::NoOp => arg,
         }
diff --git a/src/parse_trait_sig.rs b/src/parse_trait_sig.rs
index 8f578f7..be8cccb 100644
--- a/src/parse_trait_sig.rs
+++ b/src/parse_trait_sig.rs
@@ -17,8 +17,8 @@ pub enum TypeTransform {
     Into,
     Box(BoxType),
     Map(Box<TypeTransform>),
-    Collection(Box<TypeTransform>),
-    CollectionMap(Box<TypeTransform>, Box<TypeTransform>),
+    CollectionOneType(Box<TypeTransform>),
+    CollectionTwoTypes(Box<TypeTransform>, Box<TypeTransform>),
     Iterator(BoxType, Box<TypeTransform>),
     Result(Box<TypeTransform>, Box<TypeTransform>),
 }
diff --git a/src/transform.rs b/src/transform.rs
index 47e27a6..fa7877e 100644
--- a/src/transform.rs
+++ b/src/transform.rs
@@ -24,24 +24,26 @@ pub enum TransformError {
     AssocTypeInUnsupportedType(Span),
 }
 
-fn is_supported_collection(ident: &Ident) -> bool {
-    // collections added here must implement IntoIterator & FromIterator
-    // when adding a type here don't forget to document it in the README
-    ident == "Vec"
-        || ident == "VecDeque"
-        || ident == "LinkedList"
-        || ident == "HashSet"
-        || ident == "BinaryHeap"
-        || ident == "BTreeSet"
-}
+impl TypeConverter<'_> {
+    /// Returns true for types that take one generic type parameter T
+    /// and implement IntoIterator<Item=T> and FromIterator<T>
+    fn known_as_collection_with_one_type(&self, ident: &Ident) -> bool {
+        // when adding a type here don't forget to document it in the README
+        ident == "Vec"
+            || ident == "VecDeque"
+            || ident == "LinkedList"
+            || ident == "HashSet"
+            || ident == "BinaryHeap"
+            || ident == "BTreeSet"
+    }
 
-fn is_supported_collection_map(ident: &Ident) -> bool {
-    // collections added here must implement IntoIterator & FromIterator
-    // when adding a type here don't forget to document it in the README
-    ident == "HashMap" || ident == "BTreeMap"
-}
+    /// Returns true for types that take two generic type parameter K and V
+    /// and implement IntoIterator<Item=(K, V)> and FromIterator<(K, V)>
+    fn known_as_collection_with_two_types(&self, ident: &Ident) -> bool {
+        // when adding a type here don't forget to document it in the README
+        ident == "HashMap" || ident == "BTreeMap"
+    }
 
-impl TypeConverter<'_> {
     pub fn convert_type(&self, type_: &mut Type) -> Result<TypeTransform, TransformError> {
         if !iter_type(type_).any(match_assoc_type) {
             return Ok(TypeTransform::NoOp);
@@ -115,9 +117,10 @@ impl TypeConverter<'_> {
                                     || last_seg.ident == "Result"
                                 {
                                     return Ok(TypeTransform::Map(self.convert_type(arg)?.into()));
-                                } else if is_supported_collection(&last_seg.ident) && path_len == 1
+                                } else if self.known_as_collection_with_one_type(&last_seg.ident)
+                                    && path_len == 1
                                 {
-                                    return Ok(TypeTransform::Collection(
+                                    return Ok(TypeTransform::CollectionOneType(
                                         self.convert_type(arg)?.into(),
                                     ));
                                 }
@@ -136,8 +139,8 @@ impl TypeConverter<'_> {
                                         self.convert_type(arg1)?.into(),
                                         self.convert_type(arg2)?.into(),
                                     ));
-                                } else if is_supported_collection_map(&last_seg.ident) {
-                                    return Ok(TypeTransform::CollectionMap(
+                                } else if self.known_as_collection_with_two_types(&last_seg.ident) {
+                                    return Ok(TypeTransform::CollectionTwoTypes(
                                         self.convert_type(arg1)?.into(),
                                         self.convert_type(arg2)?.into(),
                                     ));