implement specialization of Base.dataids to speed-up broadcast (JuliaArrays#196)

* implement specialization of Base.dataids to speed-up broadcast

* dataids should return a tuple of UInts

Co-authored-by: Pietro Vertechi <[email protected]>

* add explicit test of dataids

Co-authored-by: Pietro Vertechi <[email protected]>

Author: ranocha

src/structarray.jl

@@ -452,3 +452,6 @@ BroadcastStyle(::Type{SA}) where SA<:StructArray = StructArrayStyle{typeof(cst(S
 
 Base.similar(bc::Broadcasted{StructArrayStyle{S}}, ::Type{ElType}) where {S<:DefaultArrayStyle,N,ElType} =
     isstructtype(ElType) ? similar(StructArray{ElType}, axes(bc)) : similar(Array{ElType}, axes(bc))
+
+# for aliasing analysis during broadcast
+Base.dataids(u::StructArray) = mapreduce(Base.dataids, (a, b) -> (a..., b...), values(components(u)), init=())

test/runtests.jl

@@ -822,6 +822,9 @@ Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray}}, ::Type{El
     @test isa(@inferred(s .+ r), StructArray)
     @test s .+ r == StructArray{ComplexF64}((s.re .+ r, s.im))
 
+    # used inside of broadcast but we also test it here explicitly
+    @test isa(@inferred(Base.dataids(s)), NTuple{N, UInt} where {N})
+
     s = StructArray{ComplexF64}((MyArray(rand(2,2)), MyArray(rand(2,2))))
     @test_throws MethodError s .+ s
 end