From ef698909223b2f29f7a6cae17f0aba3461c8803d Mon Sep 17 00:00:00 2001
From: Martin Fischer <martin@push-f.com>
Date: Mon, 22 Nov 2021 08:55:14 +0100
Subject: support HashMap and BTreeMap

---
 README.md              |  6 +++---
 src/lib.rs             |  5 +++++
 src/parse_trait_sig.rs |  1 +
 src/transform.rs       | 26 +++++++++++++++++++-------
 tests/tests.rs         |  6 +++++-
 5 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index f2a0de2..02c9d5d 100644
--- a/README.md
+++ b/README.md
@@ -91,9 +91,9 @@ Dynamize also understands if you wrap associated types in the following types:
 * `Result<_, _>`
 * `some::module::Result<_>` (type alias with fixed error type)
 * `&mut dyn Iterator<Item = _>`
-* `Vec<_>`, `VecDeque<_>`,`LinkedList<_>`
-* `HashSet<_>`, `BinaryHeap<_>`, `BTreeSet<_>`  
-  (these only work with `Into`-bounded associated types because they require `Eq`)
+* `Vec<_>`, `VecDeque<_>`, `LinkedList<_>`, `HashSet<K>`, `BinaryHeap<K>`,
+  `BTreeSet<K>`, `HashMap<K, _>`, `BTreeMap<K, _>`  
+  (for `K` only `Into`-bounded associated types work because they require `Eq`)
 
 Note that since these are resolved recursively you can actually nest these
 arbitrarily so e.g. the following also just works:
diff --git a/src/lib.rs b/src/lib.rs
index a5253d3..5e3653d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -413,6 +413,11 @@ impl TypeTransform {
                 let inner = inner.convert(quote!(x));
                 quote! {#arg.into_iter().map(|x| #inner).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::NoOp => arg,
         }
     }
diff --git a/src/parse_trait_sig.rs b/src/parse_trait_sig.rs
index c924a71..4c0b4fd 100644
--- a/src/parse_trait_sig.rs
+++ b/src/parse_trait_sig.rs
@@ -18,6 +18,7 @@ pub enum TypeTransform {
     Box(BoxType),
     Map(Box<TypeTransform>),
     Collection(Box<TypeTransform>),
+    CollectionMap(Box<TypeTransform>, Box<TypeTransform>),
     Iterator(BoxType, Box<TypeTransform>),
     Result(Box<TypeTransform>, Box<TypeTransform>),
 }
diff --git a/src/transform.rs b/src/transform.rs
index 4dc88f0..4625f4c 100644
--- a/src/transform.rs
+++ b/src/transform.rs
@@ -33,6 +33,12 @@ fn is_supported_collection(ident: &Ident) -> bool {
         || 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"
+}
+
 impl AssocTypeConversions<'_> {
     pub fn parse_type_path(&self, type_: &mut Type) -> Result<TypeTransform, TransformError> {
         if !iter_type(type_).any(match_assoc_type) {
@@ -123,14 +129,20 @@ impl AssocTypeConversions<'_> {
                         if let (GenericArgument::Type(arg1), GenericArgument::Type(arg2)) =
                             (args_iter.next().unwrap(), args_iter.next().unwrap())
                         {
-                            if (iter_type(arg1).any(match_assoc_type)
-                                || iter_type(arg2).any(match_assoc_type))
-                                && last_seg.ident == "Result"
+                            if iter_type(arg1).any(match_assoc_type)
+                                || iter_type(arg2).any(match_assoc_type)
                             {
-                                return Ok(TypeTransform::Result(
-                                    self.parse_type_path(arg1)?.into(),
-                                    self.parse_type_path(arg2)?.into(),
-                                ));
+                                if last_seg.ident == "Result" {
+                                    return Ok(TypeTransform::Result(
+                                        self.parse_type_path(arg1)?.into(),
+                                        self.parse_type_path(arg2)?.into(),
+                                    ));
+                                } else if is_supported_collection_map(&last_seg.ident) {
+                                    return Ok(TypeTransform::CollectionMap(
+                                        self.parse_type_path(arg1)?.into(),
+                                        self.parse_type_path(arg2)?.into(),
+                                    ));
+                                }
                             }
                         }
                     }
diff --git a/tests/tests.rs b/tests/tests.rs
index e004d61..925fe72 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -204,7 +204,7 @@ fn test3<T: TypeWithSuper>(some: T) {
     println!("{}", dyn_trait);
 }
 
-use std::collections::{BTreeSet, BinaryHeap, HashSet, LinkedList, VecDeque};
+use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque};
 
 #[dynamize::dynamize]
 trait Collections {
@@ -228,6 +228,10 @@ trait CollectionsRequiringEq {
     fn binary_heap(&self) -> BinaryHeap<Self::A>;
 
     fn btree_set(&self) -> BTreeSet<Self::A>;
+
+    fn hash_map(&self) -> HashMap<Self::A, Self::A>;
+
+    fn btree_map(&self) -> BTreeMap<Self::A, Self::A>;
 }
 
 #[dynamize::dynamize]
-- 
cgit v1.2.3