Store IndexName tags as a sorted symbol dictionary

src/ITensorBase.jl

@@ -23,6 +23,7 @@ include("itensor.jl")
 include("itensoroperator.jl")
 
 # `IndexName` dimname flavor and the `Index` named unit range.
+include("sorteddict.jl")
 include("index.jl")
 include("quirks.jl")
 

src/index.jl

@@ -2,9 +2,9 @@ using Accessors: @set
 using Random: AbstractRNG, RandomDevice
 using UUIDs: UUID, uuid4
 
-tagpairstring(pair::Pair) = repr(first(pair)) * "=>" * repr(last(pair))
-function tagsstring(tags::Dict{String, String})
-    tagpairs = sort(collect(tags); by = first)
+tagpairstring(pair::Pair) = string(first(pair)) * "=>" * string(last(pair))
+function tagsstring(tags)
+    tagpairs = collect(tags)  # SortedDict iterates in sorted-key order
     tagpair1, tagpair_rest = Iterators.peel(tagpairs)
     return mapreduce(*, tagpair_rest; init = tagpairstring(tagpair1)) do tagpair
         return "," * tagpairstring(tagpair)
@@ -13,14 +13,14 @@ end
 
 struct IndexName <: AbstractName
     id::UUID
-    tags::Dict{String, String}
+    tags::SortedDict{Symbol, Symbol}
     plev::Int
 end
 function IndexName(
         rng::AbstractRNG = RandomDevice(); id::UUID = uuid4(rng),
-        tags = Dict{String, String}(), plev::Int = 0
+        tags = (), plev::Int = 0
     )
-    return IndexName(id, Dict{String, String}(tags), plev)
+    return IndexName(id, to_tags(tags), plev)
 end
 function uniquename(rng::AbstractRNG, n::IndexName)
     return setid(n, uuid4(rng))
@@ -29,42 +29,77 @@ function uniquename(rng::AbstractRNG, ::Type{<:IndexName})
     return IndexName(rng)
 end
 
+to_symbol_pair(p::Pair) = Symbol(first(p)) => Symbol(last(p))
+
+# Like `Dict`, accept one or more bare `Pair`s as tags. A `Pair` iterates over
+# its two elements, so it can't fall through to the collection method below.
+to_tags(ps::Pair...) = to_tags(ps)
+to_tags(tags) = SortedDict{Symbol, Symbol}(to_symbol_pair(p) for p in tags)
+
 id(n::IndexName) = getfield(n, :id)
-tags(n::IndexName) = getfield(n, :tags)
+
+# Internal: the stored tags as `Symbol => Symbol`, used by the hot comparison,
+# hashing, and display paths. `tags` is the public string-valued view of this.
+tags_stored(n::IndexName) = getfield(n, :tags)
+
+"""
+    tags(i)
+
+Return the tags of an index or index name as an `AbstractDict` mapping tag names to
+tag values, both `AbstractString`s.
+
+The concrete dictionary type and string type are implementation details and may
+change.
+"""
+function tags(n::IndexName)
+    return SortedDict{String, String}(
+        String(k) => String(v) for (k, v) in tags_stored(n)
+    )
+end
+
 plev(n::IndexName) = getfield(n, :plev)
 
-using ConstructionBase: getfields
-Base.:(==)(n1::IndexName, n2::IndexName) = getfields(n1) == getfields(n2)
-Base.isequal(n1::IndexName, n2::IndexName) = isequal(getfields(n1), getfields(n2))
+function Base.:(==)(n1::IndexName, n2::IndexName)
+    return id(n1) == id(n2) && plev(n1) == plev(n2) && tags_stored(n1) == tags_stored(n2)
+end
+function Base.isequal(n1::IndexName, n2::IndexName)
+    return isequal(id(n1), id(n2)) &&
+        isequal(plev(n1), plev(n2)) &&
+        isequal(tags_stored(n1), tags_stored(n2))
+end
 function Base.isless(n1::IndexName, n2::IndexName)
-    t1 = (id(n1), plev(n1), keys(tags(n1)), collect(values(tags(n1))))
-    t2 = (id(n2), plev(n2), keys(tags(n2)), collect(values(tags(n2))))
+    t1 = (id(n1), plev(n1), keys(tags_stored(n1)), values(tags_stored(n1)))
+    t2 = (id(n2), plev(n2), keys(tags_stored(n2)), values(tags_stored(n2)))
     return isless(t1, t2)
 end
 function Base.hash(n::IndexName, h::UInt)
     h = hash(:IndexName, h)
     h = hash(id(n), h)
-    h = hash(tags(n), h)
     h = hash(plev(n), h)
+    h = hash(tags_stored(n), h)
     return h
 end
 
 setid(n::IndexName, id) = @set n.id = id
 settags(n::IndexName, tags) = @set n.tags = tags
 setplev(n::IndexName, plev) = @set n.plev = plev
 
-hastag(n::IndexName, tagname::String) = haskey(tags(n), tagname)
+hastag(n::IndexName, tagname) = haskey(tags_stored(n), Symbol(tagname))
 
-gettag(n::IndexName, tagname::String) = tags(n)[tagname]
-gettag(n::IndexName, tagname::String, default) = get(tags(n), tagname, default)
-function settag(n::IndexName, tagname::String, tag::String)
-    newtags = copy(tags(n))
-    newtags[tagname] = tag
+gettag(n::IndexName, tagname) = String(tags_stored(n)[Symbol(tagname)])
+function gettag(n::IndexName, tagname, default)
+    t = tags_stored(n)
+    k = Symbol(tagname)
+    return haskey(t, k) ? String(t[k]) : default
+end
+function settag(n::IndexName, tagname, tag)
+    newtags = copy(tags_stored(n))
+    newtags[Symbol(tagname)] = Symbol(tag)
     return settags(n, newtags)
 end
-function unsettag(n::IndexName, tagname::String)
-    newtags = copy(tags(n))
-    delete!(newtags, tagname)
+function unsettag(n::IndexName, tagname)
+    newtags = copy(tags_stored(n))
+    delete!(newtags, Symbol(tagname))
     return settags(n, newtags)
 end
 
@@ -122,7 +157,7 @@ shortid(id::UUID) = first(string(id), 8)
 
 function Base.show(io::IO, i::IndexName)
     idstr = "id=$(shortid(id(i)))"
-    tagsstr = !isempty(tags(i)) ? "|$(tagsstring(tags(i)))" : ""
+    tagsstr = !isempty(tags_stored(i)) ? "|$(tagsstring(tags_stored(i)))" : ""
     primestr = primestring(plev(i))
     str = "IndexName($(idstr)$(tagsstr))$(primestr)"
     print(io, str)
@@ -153,17 +188,18 @@ const Index = NamedUnitRange{IndexName}
 
 # TODO: Define for `NamedViewIndex`.
 id(i::Index) = id(name(i))
+tags_stored(i::Index) = tags_stored(name(i))
 tags(i::Index) = tags(name(i))
 plev(i::Index) = plev(name(i))
 
 # TODO: Define for `NamedViewIndex`.
-hastag(i::Index, tagname::String) = hastag(name(i), tagname)
+hastag(i::Index, tagname) = hastag(name(i), tagname)
 
 # TODO: Define for `NamedViewIndex`.
-gettag(i::Index, tagname::String) = gettag(name(i), tagname)
-gettag(i::Index, tagname::String, default) = gettag(name(i), tagname, default)
-settag(i::Index, tagname::String, tag::String) = setname(i, settag(name(i), tagname, tag))
-unsettag(i::Index, tagname::String) = setname(i, unsettag(name(i), tagname))
+gettag(i::Index, tagname) = gettag(name(i), tagname)
+gettag(i::Index, tagname, default) = gettag(name(i), tagname, default)
+settag(i::Index, tagname, tag) = setname(i, settag(name(i), tagname, tag))
+unsettag(i::Index, tagname) = setname(i, unsettag(name(i), tagname))
 
 setplev(i::Index, plev) = setname(i, setplev(name(i), plev))
 prime(i::Index) = setname(i, prime(name(i)))
@@ -185,7 +221,7 @@ end
 function Base.show(io::IO, i::Index)
     lenstr = "length=$(length(i))"
     idstr = "|id=$(shortid(id(i)))"
-    tagsstr = !isempty(tags(i)) ? "|$(tagsstring(tags(i)))" : ""
+    tagsstr = !isempty(tags_stored(i)) ? "|$(tagsstring(tags_stored(i)))" : ""
     primestr = primestring(plev(i))
     str = "Index($(lenstr)$(idstr)$(tagsstr))$(primestr)"
     print(io, str)

src/sorteddict.jl

@@ -0,0 +1,94 @@
+# Modeled on TensorKit.jl's `SortedVectorDict`, an `AbstractDict` backed by two
+# parallel vectors kept sorted by key. Comparable designs include
+# DataStructures.jl's `SortedDict` (backed by a balanced tree) and
+# Dictionaries.jl's `ArrayDictionary` (parallel arrays).
+"""
+    SortedDict{K,V} <: AbstractDict{K,V}
+
+An associative container backed by two parallel `Vector`s kept sorted by key.
+Lookup is a linear scan, which is fastest for the small key counts this is used
+for (index-name tags). Equality and hashing are structural over the sorted
+vectors, so they are cheap and order-independent by construction.
+"""
+struct SortedDict{K, V} <: AbstractDict{K, V}
+    keys::Vector{K}
+    values::Vector{V}
+    function SortedDict{K, V}(keys::Vector{K}, values::Vector{V}) where {K, V}
+        @assert length(keys) == length(values)
+        @assert issorted(keys)
+        return new{K, V}(keys, values)
+    end
+end
+
+SortedDict{K, V}() where {K, V} = SortedDict{K, V}(K[], V[])
+function SortedDict{K, V}(pairs::Vector{Pair{K, V}}) where {K, V}
+    sorted = sort(pairs; by = first)
+    return SortedDict{K, V}(map(first, sorted), map(last, sorted))
+end
+function SortedDict{K, V}(itr) where {K, V}
+    return SortedDict{K, V}(Pair{K, V}[Pair{K, V}(k, v) for (k, v) in itr])
+end
+SortedDict{K, V}(pairs::Pair...) where {K, V} = SortedDict{K, V}(collect(Pair{K, V}, pairs))
+
+Base.length(d::SortedDict) = length(getfield(d, :keys))
+Base.keys(d::SortedDict) = getfield(d, :keys)
+Base.values(d::SortedDict) = getfield(d, :values)
+
+function Base.iterate(d::SortedDict, i = 1)
+    i > length(d) && return nothing
+    return (keys(d)[i] => values(d)[i]), i + 1
+end
+
+function Base.copy(d::SortedDict{K, V}) where {K, V}
+    return SortedDict{K, V}(copy(keys(d)), copy(values(d)))
+end
+
+function Base.haskey(d::SortedDict, key)
+    # Linear scan. Use searchsortedfirst on the sorted keys for large collections.
+    return !isnothing(findfirst(isequal(key), keys(d)))
+end
+function Base.getindex(d::SortedDict, key)
+    # Linear scan. Use searchsortedfirst on the sorted keys for large collections.
+    i = findfirst(isequal(key), keys(d))
+    isnothing(i) && throw(KeyError(key))
+    return @inbounds values(d)[i]
+end
+function Base.get(d::SortedDict, key, default)
+    # Linear scan. Use searchsortedfirst on the sorted keys for large collections.
+    i = findfirst(isequal(key), keys(d))
+    return isnothing(i) ? default : @inbounds values(d)[i]
+end
+function Base.setindex!(d::SortedDict, value, key)
+    # Linear scan. Use searchsortedfirst on the sorted keys for large collections.
+    i = findfirst(isequal(key), keys(d))
+    if isnothing(i)
+        # Insertion point so the keys stay sorted. Linear scan, same upgrade as above.
+        i = something(findfirst(>(key), keys(d)), length(d) + 1)
+        insert!(keys(d), i, key)
+        insert!(values(d), i, value)
+    else
+        @inbounds values(d)[i] = value
+    end
+    return d
+end
+function Base.delete!(d::SortedDict, key)
+    # Linear scan. Use searchsortedfirst on the sorted keys for large collections.
+    i = findfirst(isequal(key), keys(d))
+    if !isnothing(i)
+        deleteat!(keys(d), i)
+        deleteat!(values(d), i)
+    end
+    return d
+end
+
+# Structural equality/hash over the (already sorted) vectors: fast and exact,
+# overriding the slower order-independent `AbstractDict` fallbacks.
+function Base.:(==)(d1::SortedDict, d2::SortedDict)
+    return keys(d1) == keys(d2) && values(d1) == values(d2)
+end
+function Base.isequal(d1::SortedDict, d2::SortedDict)
+    return isequal(keys(d1), keys(d2)) && isequal(values(d1), values(d2))
+end
+function Base.hash(d::SortedDict, h::UInt)
+    return hash(values(d), hash(keys(d), hash(:SortedDict, h)))
+end

test/test_basics.jl

@@ -28,12 +28,24 @@ using UUIDs: UUID
         @test isless(n1, n2)
         @test hash(n1) ≠ hash(n2)
 
-        for tagspec in (Dict(["X" => "Y"]), ["X" => "Y"], ("X" => "Y",), "X" => "Y")
+        for tagspec in (
+                Dict(["X" => "Y"]), ["X" => "Y"], ("X" => "Y",), "X" => "Y",
+                Dict([:X => :Y]), (:X => :Y,),
+            )
             n = IndexName(; tags = tagspec)
             @test hastag(n, "X")
+            @test hastag(n, :X)
             @test gettag(n, "X") == "Y"
+            @test gettag(n, :X) == "Y"
             @test length(tags(n)) == 1
         end
+
+        # Two-layer contract: public `tags` returns strings, internal stored layer uses Symbols.
+        n = IndexName(; tags = "X" => "Y")
+        @test tags(n) isa AbstractDict{<:AbstractString, <:AbstractString}
+        @test tags(n)["X"] == "Y"
+        @test gettag(n, "X") isa AbstractString
+        @test ITensorBase.tags_stored(n)[:X] === :Y
     end
     @testset "Index basics" begin
         i = Index(2)
@@ -135,6 +147,6 @@ using UUIDs: UUID
 
         i = settag(Index(2), "X", "Y")
         @test sprint(show, "text/plain", i) ==
-            "Index(length=2|id=$(first(string(id(i)), 8))|\"X\"=>\"Y\")"
+            "Index(length=2|id=$(first(string(id(i)), 8))|X=>Y)"
     end
 end

test/test_sorteddict.jl

@@ -0,0 +1,33 @@
+using ITensorBase: ITensorBase
+using Test: @test, @testset
+
+const SortedDict = ITensorBase.SortedDict
+
+@testset "SortedDict" begin
+    d = SortedDict{Symbol, Symbol}(:b => :y, :a => :x)
+    @test collect(keys(d)) == [:a, :b]          # sorted on construction
+    @test d[:a] == :x
+    @test haskey(d, :b)
+    @test !haskey(d, :c)
+    @test get(d, :c, :default) == :default
+    @test length(d) == 2
+
+    d[:c] = :z                                   # insert keeps sort order
+    @test collect(keys(d)) == [:a, :b, :c]
+    @test d[:c] == :z
+
+    d2 = copy(d)
+    delete!(d2, :a)
+    @test haskey(d, :a)                          # copy is independent
+    @test !haskey(d2, :a)
+
+    # order-independent equality + hash
+    e1 = SortedDict{Symbol, Symbol}(:a => :x, :b => :y)
+    e2 = SortedDict{Symbol, Symbol}(:b => :y, :a => :x)
+    @test e1 == e2
+    @test hash(e1) == hash(e2)
+    @test isequal(e1, e2)
+
+    @test length(SortedDict{Symbol, Symbol}()) == 0
+    @test sort(collect(d)) == [:a => :x, :b => :y, :c => :z]
+end