Turn off struct broadcast by default.

Author: N5N3

src/StructArrays.jl

@@ -38,5 +38,6 @@ function GPUArraysCore.backend(::Type{T}) where {T<:StructArray}
     isconsistent || throw(ArgumentError("all component arrays must have the same GPU backend"))
     return backend
 end
+_use_default_bc(::GPUArraysCore.AbstractGPUArrayStyle) = true
 
 end # module

src/staticarrays_support.jl

@@ -38,35 +38,10 @@ end
 # This looks costly, but the compiler should be able to optimize them away
 Broadcast._axes(bc::Broadcasted{<:StructStaticArrayStyle}, ::Nothing) = axes(replace_structarray(bc))
 
-to_staticstyle(@nospecialize(x::Type)) = x
-to_staticstyle(::Type{StructStaticArrayStyle{N}}) where {N} = StaticArrayStyle{N}
-
-"""
-    replace_structarray(bc::Broadcasted)
-
-An internal function transforms the `Broadcasted` with `StructArray` into
-an equivalent one without it. This is not a must if the root `BroadcastStyle`
-supports `AbstractArray`. But some `BroadcastStyle` limits the input array types, 
-e.g. `StaticArrayStyle`, thus we have to omit all `StructArray`.
-"""
-function replace_structarray(bc::Broadcasted{Style}) where {Style}
-    args = replace_structarray_args(bc.args)
-    return Broadcasted{to_staticstyle(Style)}(bc.f, args, nothing)
-end
-function replace_structarray(A::StructArray)
-    f = Instantiator(eltype(A))
-    args = Tuple(components(A))
-    return Broadcasted{StaticArrayStyle{ndims(A)}}(f, args, nothing)
-end
-replace_structarray(@nospecialize(A)) = A
-
-replace_structarray_args(args::Tuple) = (replace_structarray(args[1]), replace_structarray_args(tail(args))...)
-replace_structarray_args(::Tuple{}) = ()
-
 # StaticArrayStyle has no similar defined.
 # Overload `Base.copy` instead.
-@inline function Base.copy(bc::Broadcasted{StructStaticArrayStyle{M}}) where {M}
-    sa = copy(convert(Broadcasted{StaticArrayStyle{M}}, bc))
+@inline function try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M}
+    sa = copy(bc)
     ET = eltype(sa)
     isnonemptystructtype(ET) || return sa
     elements = Tuple(sa)

src/structarray.jl

@@ -494,7 +494,8 @@ function Base.showarg(io::IO, s::StructArray{T}, toplevel) where T
 end
 
 # broadcast
-import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, Unknown
+import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, Unknown, ArrayConflict
+using Base.Broadcast: combine_styles
 
 struct StructArrayStyle{S, N} <: AbstractArrayStyle{N} end
 
@@ -524,6 +525,49 @@ Base.@pure cst(::Type{SA}) where {SA} = combine_style_types(array_types(SA).para
 
 BroadcastStyle(::Type{SA}) where {SA<:StructArray} = StructArrayStyle{typeof(cst(SA)), ndims(SA)}()
 
+_use_default_bc(::Any) = false
+_use_default_bc(::DefaultArrayStyle) = true
+_use_default_bc(::ArrayConflict) = true
+
+function Base.copy(bc::Broadcasted{StructArrayStyle{S, N}}) where {S, N}
+    if _use_default_bc(S())
+        return @invoke copy(bc::Broadcasted)
+    else
+        return try_struct_copy(replace_structarray(bc))
+    end
+end
+try_struct_copy(bc::Broadcasted) = copy(bc) 
+
+"""
+    replace_structarray(bc::Broadcasted)
+
+An internal function transforms the `Broadcasted` with `StructArray` into
+an equivalent one without it. This is not a must if the root `BroadcastStyle`
+supports `AbstractArray`. But some `BroadcastStyle` limits the input array types, 
+e.g. `StaticArrayStyle`, thus we have to omit all `StructArray`.
+"""
+function replace_structarray(bc::Broadcasted{Style}) where {Style}
+    args = replace_structarray_args(bc.args)
+    Style′ = parant_style(Style())
+    return Broadcasted{Style′}(bc.f, args, nothing)
+end
+function replace_structarray(A::StructArray)
+    f = Instantiator(eltype(A))
+    args = Tuple(components(A))
+    Style = typeof(combine_styles(args...))
+    return Broadcasted{Style}(f, args, nothing)
+end
+replace_structarray(@nospecialize(A)) = A
+
+replace_structarray_args(args::Tuple) = (replace_structarray(args[1]), replace_structarray_args(tail(args))...)
+replace_structarray_args(::Tuple{}) = ()
+
+parant_style(@nospecialize(x)) = typeof(x)
+parant_style(::StructArrayStyle{S, N}) where {S, N} = S
+parant_style(::StructArrayStyle{S, N}) where {N, S<:AbstractArrayStyle{N}} = S
+parant_style(::StructArrayStyle{S, N}) where {S<:AbstractArrayStyle{Any}, N} = S
+parant_style(::StructArrayStyle{S, N}) where {S<:AbstractArrayStyle, N} = typeof(S(Val(N)))
+
 # Here we use `similar` defined for `S` to build the dest Array.
 function Base.similar(bc::Broadcasted{StructArrayStyle{S, N}}, ::Type{ElType}) where {S, N, ElType}
     bc′ = convert(Broadcasted{S}, bc)
@@ -532,12 +576,22 @@ end
 
 # Unwrapper to recover the behaviour defined by parent style.
 @inline function Base.copyto!(dest::AbstractArray, bc::Broadcasted{StructArrayStyle{S, N}}) where {S, N}
-    return copyto!(dest, convert(Broadcasted{S}, bc))
+    bc′ = _use_default_bc(S()) ? convert(Broadcasted{S}, bc) : replace_structarray(bc)
+    return copyto!(dest, bc′)
 end
 
 @inline function Broadcast.materialize!(::StructArrayStyle{S}, dest, bc::Broadcasted) where {S}
-    return Broadcast.materialize!(S(), dest, bc)
+    bc′ = _use_default_bc(S()) ? bc : replace_structarray(bc)
+    return Broadcast.materialize!(S(), dest, bc′)
 end
 
 # for aliasing analysis during broadcast
+function Broadcast.broadcast_unalias(dest::StructArray, src::AbstractArray)
+    if dest === src || any(Base.Fix2(===, src), components(dest))
+        return src
+    else
+        return Base.unalias(dest, src)
+    end
+end
+
 Base.dataids(u::StructArray) = mapreduce(Base.dataids, (a, b) -> (a..., b...), values(components(u)), init=())

test/runtests.jl

@@ -1183,6 +1183,7 @@ for S in (1, 2, 3)
         Base.setindex!(A::$MyArray, val, i::Int) = A.A[i] = val
         Base.size(A::$MyArray) = Base.size(A.A)
         Base.BroadcastStyle(::Type{<:$MyArray}) = Broadcast.ArrayStyle{$MyArray}()
+        StructArrays._use_default_bc(::Broadcast.ArrayStyle{$MyArray}) = true
     end
 end
 Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray1}}, ::Type{ElType}) where ElType =
@@ -1247,6 +1248,12 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
     @test Base.broadcasted(+, reshape(1:2,1,1,2), s) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{3}}
     @test Base.broadcasted(+, s, MyArray1(rand(2))) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{Any}}
 
+    # allocation test for overloaded `broadcast_unalias`
+    StructArrays._use_default_bc(::Broadcast.ArrayStyle{MyArray1}) = false
+    f(s) = s .+= 1
+    f(s)
+    @test (@allocated f(s)) == 0
+    
     # issue #185
     A = StructArray(randn(ComplexF64, 3, 3))
     B = randn(ComplexF64, 3, 3)
@@ -1317,6 +1324,14 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
         @test backend(bcmul2(sa)) === backend(sa)
         @test (sa .+= 1) === sa
     end
+
+    @testset "StructSparseArray" begin
+        a = sprand(10, 10, 0.5)
+        b = sprand(10, 10, 0.5)
+        c = StructArray{ComplexF64}((a, b))
+        d = identity.(c)
+        @test d isa SparseMatrixCSC
+    end
 end
 
 @testset "map" begin