Make map() take mixtures of StaticArray and AbstractArray

This allows StaticArray/AbstractArray linear algebra + and - to return a
static array.

Author: c42f

src/mapreduce.jl

@@ -1,25 +1,24 @@
-# Returns the common Size of the inputs (or else throws a DimensionMismatch)
-@inline same_size(a1::StaticArray, as::StaticArray...) = _same_size(Size(a1), as...)
-@inline _same_size(s::Size) = s
-@inline function _same_size(s::Size, a1::StaticArray, as::StaticArray...)
-    if s === Size(a1)
-        return _same_size(s, as...)
-    else
-        throw(DimensionMismatch("Dimensions must match. Got inputs with $s and $(Size(a1))."))
-    end
-end
-
 @inline _first(a1, as...) = a1
 
 ################
 ## map / map! ##
 ################
 
-@inline function map(f, a::StaticArray, b::StaticArray...)
-    _map(f, same_size(a, b...), a, b...)
-end
-
-@generated function _map(f, ::Size{S}, a::StaticArray...) where {S}
+@inline map(f, as::StaticArray...) =
+    _map(f, same_size(as...), as...)
+# Mixed StaticArray + AbstractArray; various versions to avoid ambiguities.
+# With the versions below, if a StaticArray isn't present in the first two
+# arguments, we'll end up in Base.map() instead.
+@inline map(f, a1::StaticArray, as::AbstractArray...) =
+    _map(f, same_size(a1, as...), a1, as...)
+@inline map(f, a1::AbstractArray, a2::StaticArray, as::AbstractArray...) =
+    _map(f, same_size(a1, a2, as...), a1, a2, as...)
+@inline map(f, a1::StaticArray, a2::AbstractArray, as::AbstractArray...) =
+    _map(f, same_size(a1, a2, as...), a1, a2, as...)
+@inline map(f, a1::StaticArray, a2::StaticArray, as::AbstractArray...) =
+    _map(f, same_size(a1, a2, as...), a1, a2, as...)
+
+@generated function _map(f, ::Size{S}, a::AbstractArray...) where {S}
     exprs = Vector{Expr}(prod(S))
     for i ∈ 1:prod(S)
         tmp = [:(a[$j][$i]) for j ∈ 1:length(a)]
@@ -29,7 +28,7 @@ end
     newT = :(Core.Inference.return_type(f, Tuple{$(eltypes...)}))
     return quote
         @_inline_meta
-        @inbounds return similar_type(typeof(_first(a...)), $newT)(tuple($(exprs...)))
+        @inbounds return similar_type(typeof(_first(a...)), $newT, Size(S))(tuple($(exprs...)))
     end
 end
 

src/traits.jl

@@ -82,11 +82,12 @@ Length(::Type{SA}) where {SA <: StaticArray} = Length(Size(SA))
 @pure length(::Size{S}) where {S} = length(S)
 @pure length_val{S}(::Size{S}) = Val{length(S)}
 
-@pure Base.:(==){S}(::Size{S}, s::Tuple{Vararg{Int}}) = S == s
-@pure Base.:(==){S}(s::Tuple{Vararg{Int}}, ::Size{S}) = s == S
+# Note - using === here, as Base doesn't inline == for tuples as of julia-0.6
+@pure Base.:(==){S}(::Size{S}, s::Tuple{Vararg{Int}}) = S === s
+@pure Base.:(==){S}(s::Tuple{Vararg{Int}}, ::Size{S}) = s === S
 
-@pure Base.:(!=){S}(::Size{S}, s::Tuple{Vararg{Int}}) = S != s
-@pure Base.:(!=){S}(s::Tuple{Vararg{Int}}, ::Size{S}) = s != S
+@pure Base.:(!=){S}(::Size{S}, s::Tuple{Vararg{Int}}) = S !== s
+@pure Base.:(!=){S}(s::Tuple{Vararg{Int}}, ::Size{S}) = s !== S
 
 @pure Base.prod{S}(::Size{S}) = prod(S)
 
@@ -104,3 +105,32 @@ Length(::Type{SA}) where {SA <: StaticArray} = Length(Size(SA))
 
 # The generated functions work with length, etc...
 @propagate_inbounds unroll_tuple(f, ::Length{L}) where {L} = unroll_tuple(f, Val{L})
+
+
+"""
+Static or runtime size of an array
+"""
+const SRSize = Union{Size,Tuple{Vararg{Int}}}
+
+"""
+Return either the statically known Size() or runtime size()
+"""
+@inline _size(a) = size(a)
+@inline _size(a::StaticArray) = Size(a)
+
+# Return first static Size from a set of arrays
+@inline _first_static_size(a1::StaticArray, as...) = Size(a1)
+@inline _first_static_size(a1, as...) = _first_static_size(as...)
+@inline _first_static_size() = throw(ArgumentError("No StaticArray found in argument list"))
+
+# Returns the common Size of the inputs (or else throws a DimensionMismatch)
+@inline same_size(as...) = _same_size(_first_static_size(as...), as...)
+@inline function _same_size(s::SRSize, a1, as...)
+    if s == _size(a1)
+        return _same_size(s, as...)
+    else
+        throw(DimensionMismatch("Dimensions must match. Got inputs with $s and $(_size(a1))."))
+    end
+end
+@inline _same_size(s::SRSize) = s
+

test/linalg.jl

@@ -17,11 +17,10 @@
         v3 = [2,4,6,8]
         v4 = [4,3,2,1]
 
-        # We broke "inferrable" sizes of AbstractVectors for vector+vector, matrix*vector, etc...
-        @test_broken @inferred(v1 + v4) === @SVector [6, 7, 8, 9]
-        @test_broken @inferred(v3 + v2) === @SVector [6, 7, 8, 9]
-        @test_broken @inferred(v1 - v4) === @SVector [-2, 1, 4, 7]
-        @test_broken @inferred(v3 - v2) === @SVector [-2, 1, 4, 7]
+        @test @inferred(v1 + v4) === @SVector [6, 7, 8, 9]
+        @test @inferred(v3 + v2) === @SVector [6, 7, 8, 9]
+        @test @inferred(v1 - v4) === @SVector [-2, 1, 4, 7]
+        @test @inferred(v3 - v2) === @SVector [-2, 1, 4, 7]
     end
 
     @testset "Interaction with `UniformScaling`" begin

test/mapreduce.jl

@@ -9,8 +9,8 @@
 
         @test @inferred(map(-, v1)) === @SVector [-2, -4, -6, -8]
         @test @inferred(map(+, v1, v2)) === @SVector [6, 7, 8, 9]
-        @test_broken @inferred(map(+, normal_v1, v2)) === @SVector [6, 7, 8, 9]
-        @test_broken @inferred(map(+, v1, normal_v2)) === @SVector [6, 7, 8, 9]
+        @test @inferred(map(+, normal_v1, v2)) === @SVector [6, 7, 8, 9]
+        @test @inferred(map(+, v1, normal_v2)) === @SVector [6, 7, 8, 9]
 
         map!(+, mv, v1, v2)
         @test mv == @MVector [6, 7, 8, 9]