support boxed Iterator in Option

5 files changed, 26 insertions, 14 deletions

diff --git a/src/lib.rs b/src/lib.rs
index fc81e42..1f75023 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -390,9 +390,9 @@ impl TypeTransform {
                 let inner = opt.convert(quote!(x));
                 quote! {#arg.map(|x| #inner)}
             }
-            TypeTransform::Iterator(inner) => {
+            TypeTransform::Iterator(box_type, inner) => {
                 let inner = inner.convert(quote!(x));
-                quote! {Box::new(#arg.map(|x| #inner))}
+                quote! {Box::new(#arg.map(|x| #inner)) as #box_type}
             }
             TypeTransform::Result(ok, err) => {
                 let map_ok = !matches!(ok.as_ref(), TypeTransform::NoOp);
diff --git a/src/parse_assoc_type.rs b/src/parse_assoc_type.rs
index 85f3723..048e58c 100644
--- a/src/parse_assoc_type.rs
+++ b/src/parse_assoc_type.rs
@@ -1,7 +1,7 @@
-use proc_macro2::Span;
+use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens};
 use syn::spanned::Spanned;
-use syn::{GenericArgument, Ident, Path, PathArguments, PathSegment, TraitItemType, Type};
+use syn::{GenericArgument, Ident, PathArguments, PathSegment, TraitItemType, Type};
 
 use crate::match_assoc_type;
 use crate::parse_trait_sig::TypeTransform;
@@ -16,16 +16,16 @@ pub enum AssocTypeParseError {
 
 #[derive(Debug, Clone)]
 pub struct BoxType {
-    trait_name: Path,
-    placeholder_lifetime: bool,
+    pub inner: TokenStream,
+    pub placeholder_lifetime: bool,
 }
 
 impl ToTokens for BoxType {
     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
-        let path = &self.trait_name;
+        let inner = &self.inner;
         match self.placeholder_lifetime {
-            true => tokens.extend(quote! {Box<dyn #path + '_>}),
-            false => tokens.extend(quote! {Box<dyn #path>}),
+            true => tokens.extend(quote! {Box<#inner + '_>}),
+            false => tokens.extend(quote! {Box<#inner>}),
         }
     }
 }
@@ -79,10 +79,11 @@ pub fn parse_assoc_type(
                 }
             }
         }
+        let path = &bound.path;
         return Ok((
             &assoc_type.ident,
             DestType::Box(BoxType {
-                trait_name: bound.path.clone(),
+                inner: quote! {dyn #path},
                 placeholder_lifetime: !lifetime_bounds(&assoc_type.bounds)
                     .any(|l| l.ident == "static"),
             }),
diff --git a/src/parse_trait_sig.rs b/src/parse_trait_sig.rs
index cf2d7da..ab2c437 100644
--- a/src/parse_trait_sig.rs
+++ b/src/parse_trait_sig.rs
@@ -17,7 +17,7 @@ pub enum TypeTransform {
     Into,
     Box(BoxType),
     Map(Box<TypeTransform>),
-    Iterator(Box<TypeTransform>),
+    Iterator(BoxType, Box<TypeTransform>),
     Result(Box<TypeTransform>, Box<TypeTransform>),
 }
 
diff --git a/src/transform.rs b/src/transform.rs
index fff9c6c..69929cc 100644
--- a/src/transform.rs
+++ b/src/transform.rs
@@ -9,7 +9,7 @@ use syn::{
 
 use crate::{
     filter_map_assoc_paths, match_assoc_type,
-    parse_assoc_type::DestType,
+    parse_assoc_type::{BoxType, DestType},
     parse_trait_sig::{MethodParseError, TypeTransform},
     syn_utils::{iter_path, iter_type},
 };
@@ -56,8 +56,15 @@ impl AssocTypeConversions<'_> {
                                             {
                                                 let inner =
                                                     self.parse_type_path(&mut binding.ty)?;
-                                                *type_ = Type::Verbatim(quote! {Box<#elem + '_>});
-                                                return Ok(TypeTransform::Iterator(inner.into()));
+                                                let box_type = BoxType {
+                                                    inner: quote! {#elem},
+                                                    placeholder_lifetime: true,
+                                                };
+                                                *type_ = Type::Verbatim(quote! {#box_type});
+                                                return Ok(TypeTransform::Iterator(
+                                                    box_type,
+                                                    inner.into(),
+                                                ));
                                             }
                                         }
                                     }
diff --git a/tests/tests.rs b/tests/tests.rs
index 8bbb5b3..b326b4b 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -216,4 +216,8 @@ trait FunIter {
     type A: std::error::Error;
 
     fn foobar<F: Fn(&mut dyn Iterator<Item = Self::A>)>(&mut self, f: F);
+
+    fn foobar1<G: Fn(Option<&mut dyn Iterator<Item = Self::A>>)>(&mut self, f: G);
+
+    fn foobar2<H: Fn(&mut dyn Iterator<Item = &mut dyn Iterator<Item = Self::A>>)>(&mut self, f: H);
 }