Make instantiate generate static axes

add `static_combine_axes` to handle Tuple correctly

Author: N5N3

src/broadcast.jl

@@ -4,6 +4,7 @@
 
 import Base.Broadcast:
 BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, materialize!
+import Base.Broadcast: combine_axes, instantiate, _broadcast_getindex, broadcast_shape, Style
 import Base.Broadcast: _bcs1  # for SOneTo axis information
 using Base.Broadcast: _bcsm
 # Add a new BroadcastStyle for StaticArrays, derived from AbstractArrayStyle
@@ -19,6 +20,42 @@ BroadcastStyle(::StaticArrayStyle{M}, ::DefaultArrayStyle{N}) where {M,N} =
     DefaultArrayStyle(Val(max(M, N)))
 BroadcastStyle(::StaticArrayStyle{M}, ::DefaultArrayStyle{0}) where {M} =
     StaticArrayStyle{M}()
+
+# combine_axes overload (for Tuple)
+@inline static_combine_axes(A, B...) = broadcast_shape(static_axes(A), static_combine_axes(B...))
+static_combine_axes(A) = static_axes(A)
+static_axes(A) = axes(A)
+static_axes(x::Tuple) = (SOneTo{length(x)}(),)
+static_axes(bc::Broadcasted{Style{Tuple}}) = static_combine_axes(bc.args...)
+Broadcast._axes(bc::Broadcasted{<:StaticArrayStyle}, ::Nothing) = static_combine_axes(bc.args...)
+
+# instantiate overload
+@inline function instantiate(B::Broadcasted{StaticArrayStyle{M}}) where M
+    if B.axes isa Tuple{Vararg{SOneTo}} || B.axes isa Tuple && length(B.axes) > M
+        return invoke(instantiate, Tuple{Broadcasted}, B)
+    elseif B.axes isa Nothing
+        ax = static_combine_axes(B.args...)
+        return Broadcasted{StaticArrayStyle{M}}(B.f, B.args, ax)
+    else
+        # We need to update B.axes for `broadcast!` if it's not static and `ndims(dest) < M`.
+        ax = static_check_broadcast_shape(B.axes, static_combine_axes(B.args...))
+        return Broadcasted{StaticArrayStyle{M}}(B.f, B.args, ax)
+    end
+end
+@inline function static_check_broadcast_shape(shp::Tuple, Ashp::Tuple{Vararg{SOneTo}})
+    ax1 = if length(Ashp[1]) == 1
+        shp[1]
+    elseif Ashp[1] == shp[1]
+        Ashp[1]
+    else
+        throw(DimensionMismatch("array could not be broadcast to match destination"))
+    end
+    (ax1, static_check_broadcast_shape(Base.tail(shp), Base.tail(Ashp))...)
+end
+static_check_broadcast_shape(::Tuple{}, ::Tuple{SOneTo,Vararg{SOneTo}}) =
+    throw(DimensionMismatch("cannot broadcast array to have fewer non-singleton dimensions"))
+static_check_broadcast_shape(::Tuple{}, ::Tuple{SOneTo{1},Vararg{SOneTo{1}}}) = ()
+static_check_broadcast_shape(::Tuple{}, ::Tuple{}) = ()
 # copy overload
 @inline function Base.copy(B::Broadcasted{StaticArrayStyle{M}}) where M
     flat = Broadcast.flatten(B); as = flat.args; f = flat.f
@@ -42,8 +79,14 @@ end
 
 # Resolving priority between dynamic and static axes
 _bcs1(a::SOneTo, b::SOneTo) = _bcsm(b, a) ? b : (_bcsm(a, b) ? a : throw(DimensionMismatch("arrays could not be broadcast to a common size")))
-_bcs1(a::SOneTo, b::Base.OneTo) = _bcs1(Base.OneTo(a), b)
-_bcs1(a::Base.OneTo, b::SOneTo) = _bcs1(a, Base.OneTo(b))
+function _bcs1(a::SOneTo, b::Base.OneTo)
+    length(a) == 1 && return b
+    if length(b) != length(a) && length(b) != 1
+        throw(DimensionMismatch("arrays could not be broadcast to a common size"))
+    end
+    return a
+end
+_bcs1(a::Base.OneTo, b::SOneTo) = _bcs1(b, a)
 
 ###################################################
 ## Internal broadcast machinery for StaticArrays ##