Merge branch mixed-sarray-array-mapping

Author: Chris Foster

src/mapreduce.jl

@@ -1,25 +1,16 @@
-# 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...)
+# The following type signature for map() matches any list of AbstractArrays,
+# provided at least one is a static array.
+@inline function map(f, as::Union{SA,AbstractArray}...) where {SA<:StaticArray}
+    _map(f, same_size(as...), as...)
 end
 
-@generated function _map(f, ::Size{S}, a::StaticArray...) where {S}
+@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 +20,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,29 @@ 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})
+
+
+"""
+Return either the statically known Size() or runtime size()
+"""
+@inline _size(a) = size(a)
+@inline _size(a::StaticArray) = Size(a)
+
+# Return static array from a set of arrays
+@inline _first_static(a1::StaticArray, as...) = a1
+@inline _first_static(a1, as...) = _first_static(as...)
+@inline _first_static() = throw(ArgumentError("No StaticArray found in argument list"))
+
+"""
+Returns the common Size of the inputs (or else throws a DimensionMismatch)
+"""
+@inline function same_size(as...)
+    s = Size(_first_static(as...))
+    _sizes_match(s, as...) || _throw_size_mismatch(as...)
+    s
+end
+@inline _sizes_match(s::Size, a1, as...) = ((s == _size(a1)) ? _sizes_match(s, as...) : false)
+@inline _sizes_match(s::Size) = true
+@noinline function _throw_size_mismatch(as...)
+    throw(DimensionMismatch("Sizes $(map(_size, as)) of input arrays do not match"))
+end

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]