Fix NaiveOrderedSet broadcasting mixed with Tuples and arrays (#34)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "NamedDimsArrays"
 uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/abstractnameddimsarray.jl

@@ -175,7 +175,7 @@ function combine_nameddimsarraytype(
 end
 combine_nameddimsarraytype(::Type{T}, ::Type{T}) where {T<:AbstractNamedDimsArray} = T
 
-using Base.Broadcast: Broadcasted, Style
+using Base.Broadcast: AbstractArrayStyle, Broadcasted, Style
 
 struct NaiveOrderedSet{Values}
   values::Values
@@ -188,10 +188,15 @@ Base.keys(s::NaiveOrderedSet) = Base.OneTo(length(s))
 Base.:(==)(s1::NaiveOrderedSet, s2::NaiveOrderedSet) = issetequal(values(s1), values(s2))
 Base.iterate(s::NaiveOrderedSet, args...) = iterate(values(s), args...)
 Base.getindex(s::NaiveOrderedSet, I::Int) = values(s)[I]
+# TODO: Required in Julia 1.10, delete when we drop support for that.
+Base.getindex(s::NaiveOrderedSet, I::CartesianIndex{1}) = values(s)[I]
 Base.get(s::NaiveOrderedSet, I::Integer, default) = get(values(s), I, default)
 Base.invperm(s::NaiveOrderedSet) = NaiveOrderedSet(invperm(values(s)))
 Base.Broadcast._axes(::Broadcasted, axes::NaiveOrderedSet) = axes
 Base.Broadcast.BroadcastStyle(::Type{<:NaiveOrderedSet}) = Style{NaiveOrderedSet}()
+Base.Broadcast.BroadcastStyle(::Style{Tuple}, ::Style{NaiveOrderedSet}) = Style{Tuple}()
+Base.Broadcast.BroadcastStyle(s1::AbstractArrayStyle{0}, s2::Style{NaiveOrderedSet}) = s2
+Base.Broadcast.BroadcastStyle(s1::AbstractArrayStyle, s2::Style{NaiveOrderedSet}) = s1
 Base.Broadcast.broadcastable(s::NaiveOrderedSet) = s
 Base.to_shape(s::NaiveOrderedSet) = s
 

test/test_basics.jl

@@ -292,7 +292,13 @@ using Test: @test, @test_throws, @testset
     end
   end
   @testset "NaiveOrderedSet" begin
+    # Broadcasting
+    s = NaiveOrderedSet((1, 2))
+    @test s .+ [3, 4] == [4, 6]
+    @test s .+ (3, 4) === (4, 6)
+
     s = NaiveOrderedSet(("a", "b", "c"))
+    @test all(s .== ("a", "b", "c"))
     @test values(s) == ("a", "b", "c")
     @test Tuple(s) == ("a", "b", "c")
     @test s[1] == "a"