Improve inferability of resursive map on nested StaticArrays (#594)

Remove the `::Size` argument to `_map`, instead determining the size  
internally.

Author: martinholters

src/mapreduce.jl

@@ -8,16 +8,27 @@
 # `map(f, as::Union{<:StaticArray,AbstractArray}...)` which included at least one `StaticArray`
 # this is not the case on 0.7 and we instead hope to find a StaticArray in the first two arguments.
 @inline function map(f, a1::StaticArray, as::AbstractArray...)
-    _map(f, same_size(a1, as...), a1, as...)
+    _map(f, a1, as...)
 end
 @inline function map(f, a1::AbstractArray, a2::StaticArray, as::AbstractArray...)
-    _map(f, same_size(a1, a2, as...), a1, a2, as...)
+    _map(f, a1, a2, as...)
 end
 @inline function map(f, a1::StaticArray, a2::StaticArray, as::AbstractArray...)
-    _map(f, same_size(a1, a2, as...), a1, a2, as...)
+    _map(f, a1, a2, as...)
 end
 
-@generated function _map(f, ::Size{S}, a::AbstractArray...) where {S}
+@generated function _map(f, a::AbstractArray...)
+    i = findfirst(ai -> ai <: StaticArray, a)
+    if i === nothing
+        return :(throw(ArgumentError("No StaticArray found in argument list")))
+    end
+    # Passing the Size as an argument to _map leads to inference issues when
+    # recursively mapping over nested StaticArrays (see issue #593). Calling
+    # Size in the generator here is valid because a[i] is known to be a
+    # StaticArray for which the default Size method is correct. If wrapped
+    # StaticArrays (with a custom Size method) are to be supported, this will
+    # no longer be valid.
+    S = Size(a[i])
     exprs = Vector{Expr}(undef, prod(S))
     for i ∈ 1:prod(S)
         tmp = [:(a[$j][$i]) for j ∈ 1:length(a)]
@@ -26,8 +37,9 @@ end
 
     return quote
         @_inline_meta
+        S = same_size(a...)
         @inbounds elements = tuple($(exprs...))
-        @inbounds return similar_type(typeof(_first(a...)), eltype(elements), Size(S))(elements)
+        @inbounds return similar_type(typeof(_first(a...)), eltype(elements), S)(elements)
     end
 end
 

test/mapreduce.jl

@@ -149,4 +149,11 @@ end
 
         @test ((@SVector Int64[]) + (@SVector Int64[])) == (@SVector Int64[])
     end
+    @testset "Nested SVectors" begin
+        # issue #593
+        v = SVector(SVector(3, 2), SVector(5, 7))
+        @test @inferred(v + v) == SVector(SVector(6, 4), SVector(10, 14))
+        v = SVector(SVector(3, 2, 1), SVector(5, 7, 9))
+        @test @inferred(v + v) == SVector(SVector(6, 4, 2), SVector(10, 14, 18))
+    end
 end