Refactor to generalize map() and improve codegen for same_size()

Author: Chris Foster

src/mapreduce.jl

@@ -1,22 +1,10 @@
-@inline _first(a1, as...) = a1
-
 ################
 ## map / map! ##
 ################
 
-@inline map(f, as::StaticArray...) =
+@inline function map(f, as::Union{SA,AbstractArray}...) where {SA<: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...)
+end
 
 @generated function _map(f, ::Size{S}, a::AbstractArray...) where {S}
     exprs = Vector{Expr}(prod(S))
@@ -28,7 +16,7 @@
     newT = :(Core.Inference.return_type(f, Tuple{$(eltypes...)}))
     return quote
         @_inline_meta
-        @inbounds return similar_type(typeof(_first(a...)), $newT, Size(S))(tuple($(exprs...)))
+        @inbounds return similar_type(typeof(_first_static(a...)), $newT, Size(S))(tuple($(exprs...)))
     end
 end
 

src/traits.jl

@@ -113,19 +113,21 @@ 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::Size, 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::Size) = s
+# 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