Drop combine_size

1. `axes` should be static.
2. The implementation for broadcast(1) are unified)
3. simplify `broadcast_index`

Author: N5N3

src/abstractarray.jl

@@ -14,6 +14,7 @@ Base.axes(s::StaticArray) = _axes(Size(s))
 end
 Base.axes(rv::Adjoint{<:Any,<:StaticVector})   = (SOneTo(1), axes(rv.parent)...)
 Base.axes(rv::Transpose{<:Any,<:StaticVector}) = (SOneTo(1), axes(rv.parent)...)
+Base.axes(d::Diagonal{<:Any,<:StaticVector}) = (ax = axes(d.diag, 1); (ax, ax))
 
 Base.eachindex(::IndexLinear, a::StaticArray) = SOneTo(length(a))
 

src/broadcast.jl

@@ -60,21 +60,25 @@ static_check_broadcast_shape(::Tuple{}, ::Tuple{}) = ()
 @inline function Base.copy(B::Broadcasted{StaticArrayStyle{M}}) where M
     flat = Broadcast.flatten(B); as = flat.args; f = flat.f
     argsizes = broadcast_sizes(as...)
-    destsize = combine_sizes(argsizes)
-    _broadcast(f, destsize, argsizes, as...)
+    ax = axes(B)
+    if ax isa Tuple{Vararg{SOneTo}}
+        return _broadcast(f, Size(map(length, ax)), argsizes, as...)
+    end
+    return copy(convert(Broadcasted{DefaultArrayStyle{M}}, B))
 end
 # copyto! overloads
 @inline Base.copyto!(dest, B::Broadcasted{<:StaticArrayStyle}) = _copyto!(dest, B)
 @inline Base.copyto!(dest::AbstractArray, B::Broadcasted{<:StaticArrayStyle}) = _copyto!(dest, B)
 @inline function _copyto!(dest, B::Broadcasted{StaticArrayStyle{M}}) where M
     flat = Broadcast.flatten(B); as = flat.args; f = flat.f
     argsizes = broadcast_sizes(as...)
-    destsize = combine_sizes((Size(dest), argsizes...))
-    if Length(destsize) === Length{Dynamic()}()
-        # destination dimension cannot be determined statically; fall back to generic broadcast!
-        return copyto!(dest, convert(Broadcasted{DefaultArrayStyle{M}}, B))
+    ax = axes(B)
+    if ax isa Tuple{Vararg{SOneTo}}
+        @boundscheck axes(dest) == ax || Broadcast.throwdm(axes(dest), ax)
+        return _broadcast!(f, Size(map(length, ax)), dest, argsizes, as...)
     end
-    _broadcast!(f, destsize, dest, argsizes, as...)
+    # destination dimension cannot be determined statically; fall back to generic broadcast!
+    return copyto!(dest, convert(Broadcasted{DefaultArrayStyle{M}}, B))
 end
 
 # Resolving priority between dynamic and static axes
@@ -99,45 +103,13 @@ _bcs1(a::Base.OneTo, b::SOneTo) = _bcs1(b, a)
 @inline broadcast_size(a::AbstractArray) = Size(a)
 @inline broadcast_size(a::Tuple) = Size(length(a))
 
-function broadcasted_index(oldsize, newindex)
-    index = ones(Int, length(oldsize))
-    for i = 1:length(oldsize)
-        if oldsize[i] != 1
-            index[i] = newindex[i]
-        end
-    end
-    return LinearIndices(oldsize)[index...]
-end
-
-# similar to Base.Broadcast.combine_indices:
-@generated function combine_sizes(s::Tuple{Vararg{Size}})
-    sizes = [sz.parameters[1] for sz ∈ s.parameters]
-    ndims = 0
-    for i = 1:length(sizes)
-        ndims = max(ndims, length(sizes[i]))
-    end
-    newsize = StaticDimension[Dynamic() for _ = 1 : ndims]
-    for i = 1:length(sizes)
-        s = sizes[i]
-        for j = 1:length(s)
-            if s[j] isa Dynamic
-                continue
-            elseif newsize[j] isa Dynamic || newsize[j] == 1
-                newsize[j] = s[j]
-            elseif newsize[j] ≠ s[j] && s[j] ≠ 1
-                throw(DimensionMismatch("Tried to broadcast on inputs sized $sizes"))
-            end
-        end
-    end
-    quote
-        @_inline_meta
-        Size($(tuple(newsize...)))
-    end
+broadcast_getindex(::Tuple{}, i::Int, I::CartesianIndex) = return :(_broadcast_getindex(a[$i], $I))
+function broadcast_getindex(oldsize::Tuple, i::Int, newindex::CartesianIndex)
+    li = LinearIndices(oldsize)
+    ind = _broadcast_getindex(li, newindex)
+    return :(a[$i][$ind])
 end
 
-scalar_getindex(x) = x
-scalar_getindex(x::Ref) = x[]
-
 isstatic(::StaticArray) = true
 isstatic(::Transpose{<:Any, <:StaticArray}) = true
 isstatic(::Adjoint{<:Any, <:StaticArray}) = true
@@ -161,13 +133,11 @@ end
 
 @generated function __broadcast(f, ::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where newsize
     sizes = [sz.parameters[1] for sz ∈ s.parameters]
+
     indices = CartesianIndices(newsize)
     exprs = similar(indices, Expr)
     for (j, current_ind) ∈ enumerate(indices)
-        exprs_vals = [
-            (!(a[i] <: AbstractArray || a[i] <: Tuple) ? :(scalar_getindex(a[$i])) : :(a[$i][$(broadcasted_index(sizes[i], current_ind))]))
-            for i = 1:length(sizes)
-        ]
+        exprs_vals = (broadcast_getindex(sz, i, current_ind) for (i, sz) in enumerate(sizes))
         exprs[j] = :(f($(exprs_vals...)))
     end
 
@@ -181,27 +151,18 @@ end
 ## Internal broadcast! machinery for StaticArrays ##
 ####################################################
 
-@generated function _broadcast!(f, ::Size{newsize}, dest::AbstractArray, s::Tuple{Vararg{Size}}, as...) where {newsize}
-    sizes = [sz.parameters[1] for sz ∈ s.parameters]
-    sizes = tuple(sizes...)
-
-    # TODO: this could also be done outside the generated function:
-    sizematch(Size{newsize}(), Size(dest)) ||
-        throw(DimensionMismatch("Tried to broadcast to destination sized $newsize from inputs sized $sizes"))
+@generated function _broadcast!(f, ::Size{newsize}, dest::AbstractArray, s::Tuple{Vararg{Size}}, a...) where {newsize}
+    sizes = [sz.parameters[1] for sz in s.parameters]
 
     indices = CartesianIndices(newsize)
     exprs = similar(indices, Expr)
     for (j, current_ind) ∈ enumerate(indices)
-        exprs_vals = [
-            (!(as[i] <: AbstractArray || as[i] <: Tuple) ? :(as[$i][]) : :(as[$i][$(broadcasted_index(sizes[i], current_ind))]))
-            for i = 1:length(sizes)
-        ]
+        exprs_vals = (broadcast_getindex(sz, i, current_ind) for (i, sz) in enumerate(sizes))
         exprs[j] = :(dest[$j] = f($(exprs_vals...)))
     end
 
     return quote
-        @_propagate_inbounds_meta
-        @boundscheck sizematch($(Size{newsize}()), dest) || throw(DimensionMismatch("array could not be broadcast to match destination"))
+        @_inline_meta
         @inbounds $(Expr(:block, exprs...))
         return dest
     end

src/precompile.jl

@@ -24,6 +24,10 @@ function _precompile_()
     # Some expensive generators
     @assert precompile(Tuple{typeof(which(__broadcast,(Any,Size,Tuple{Vararg{Size}},Vararg{Any},)).generator.gen),Any,Any,Any,Any,Any,Any})
     @assert precompile(Tuple{typeof(which(_zeros,(Size,Type{<:StaticArray},)).generator.gen),Any,Any,Any,Type,Any})
-    @assert precompile(Tuple{typeof(which(combine_sizes,(Tuple{Vararg{Size}},)).generator.gen),Any,Any})
     @assert precompile(Tuple{typeof(which(_mapfoldl,(Any,Any,Colon,Any,Size,Vararg{StaticArray},)).generator.gen),Any,Any,Any,Any,Any,Any,Any,Any})
+
+    # broadcast_getindex
+    for m = 0:5, n = m:5 
+        @assert precompile(Tuple{typeof(broadcast_getindex),NTuple{m,Int},Int,CartesianIndex{n}})
+    end
 end