Handle mixed size iterators

Author: mschauer

src/flatten.jl

@@ -1,6 +1,6 @@
 # Special case flatten of iterators of static arrays.
 import Base.Iterators: flatten_iteratorsize, flatten_length
-flatten_iteratorsize(::Union{Base.HasShape, Base.HasLength}, ::Type{<:StaticArray}) = Base.HasLength()
-function flatten_length(f, T::Type{<:StaticArray})
+flatten_iteratorsize(::Union{Base.HasShape, Base.HasLength}, ::Type{<:StaticArray{S}}) where {S} = Base.HasLength()
+function flatten_length(f, T::Type{<:StaticArray{S}}) where {S}
     length(T)*length(f.it)
 end

test/flatten.jl

@@ -10,4 +10,6 @@ using StaticArrays, Test
         @test collect(Iterators.flatten(typeof(x)[])) == []
         @test collect(Iterators.flatten(X)) == [x..., x..., x...]
     end
+    @test collect(Iterators.flatten([SVector(1,1), SVector(1)])) == [1,1,1]
+    @test_throws ArgumentError length(Iterators.flatten([SVector(1,1), SVector(1)]))
 end