From 67585f872ab3b640916094b41f907ff0b705700f Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 15:19:21 -0400 Subject: [PATCH 1/4] Add internal SortedDict for sorted, vector-backed associative storage A small `AbstractDict` backed by two parallel sorted vectors with linear-scan lookup, suited to the small key counts of index-name tags. Equality and hashing are structural over the sorted vectors, so they are cheap and order-independent by construction. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/ITensorBase.jl | 1 + src/sorteddict.jl | 100 ++++++++++++++++++++++++++++++++++++++++ test/test_sorteddict.jl | 33 +++++++++++++ 3 files changed, 134 insertions(+) create mode 100644 src/sorteddict.jl create mode 100644 test/test_sorteddict.jl diff --git a/src/ITensorBase.jl b/src/ITensorBase.jl index 278bfa1b..74499cae 100644 --- a/src/ITensorBase.jl +++ b/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") diff --git a/src/sorteddict.jl b/src/sorteddict.jl new file mode 100644 index 00000000..76673f3c --- /dev/null +++ b/src/sorteddict.jl @@ -0,0 +1,100 @@ +""" + 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 + +# Linear scan: index of the first key not less than `key`. Fastest for small +# `length(d)`; swap in `searchsortedfirst` if a large-`n` use case appears. +function _searchsortedfirst(ks::Vector, key) + i = 1 + @inbounds while i <= length(ks) && isless(ks[i], key) + i += 1 + end + return i +end + +function Base.haskey(d::SortedDict{K}, key) where {K} + k = convert(K, key) + i = _searchsortedfirst(keys(d), k) + return i <= length(d) && isequal(keys(d)[i], k) +end +function Base.getindex(d::SortedDict{K}, key) where {K} + k = convert(K, key) + i = _searchsortedfirst(keys(d), k) + (i <= length(d) && isequal(keys(d)[i], k)) || throw(KeyError(key)) + return @inbounds values(d)[i] +end +function Base.get(d::SortedDict{K}, key, default) where {K} + k = convert(K, key) + i = _searchsortedfirst(keys(d), k) + return (i <= length(d) && isequal(keys(d)[i], k)) ? (@inbounds values(d)[i]) : default +end +function Base.setindex!(d::SortedDict{K, V}, value, key) where {K, V} + k = convert(K, key) + v = convert(V, value) + i = _searchsortedfirst(keys(d), k) + if i <= length(d) && isequal(keys(d)[i], k) + @inbounds values(d)[i] = v + else + insert!(keys(d), i, k) + insert!(values(d), i, v) + end + return d +end +function Base.delete!(d::SortedDict{K}, key) where {K} + k = convert(K, key) + i = _searchsortedfirst(keys(d), k) + if i <= length(d) && isequal(keys(d)[i], k) + 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 diff --git a/test/test_sorteddict.jl b/test/test_sorteddict.jl new file mode 100644 index 00000000..24df746c --- /dev/null +++ b/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 From 1640b07b890c046e9baa876550b067bfe83ea4a7 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 15:19:33 -0400 Subject: [PATCH 2/4] Store IndexName tags as a sorted symbol dictionary Back IndexName tags with the internal SortedDict{Symbol,Symbol} instead of a Dict{String,String}. Symbol keys and values compare and hash faster than strings, and the sorted backing makes equality, hashing, and ordering plain vector operations with no per-call allocation. IndexName comparison now short-circuits in id, then plev, then tags, so distinct ids never touch the tags and comparing an index against a primed or retagged copy of itself rejects on the prime level instead of walking the full tag set. The public tag API stays string-valued: gettag returns a String, tags returns an AbstractDict from strings to strings, and settag/hastag and the IndexName constructor accept either strings or symbols. An internal tags_stored exposes the raw symbol dictionary for the comparison, hashing, and display paths. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/index.jl | 97 +++++++++++++++++++++++++++++++-------------- test/test_basics.jl | 16 +++++++- 2 files changed, 81 insertions(+), 32 deletions(-) diff --git a/src/index.jl b/src/index.jl index 0d2dd777..32d7b1c2 100644 --- a/src/index.jl +++ b/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,78 @@ function uniquename(rng::AbstractRNG, ::Type{<:IndexName}) return IndexName(rng) end +to_tags(tags::SortedDict{Symbol, Symbol}) = tags +to_tags(tag::Pair) = to_tags((tag,)) +function to_tags(tags) + return SortedDict{Symbol, Symbol}( + Symbol(first(p)) => Symbol(last(p)) for p in tags + ) +end + 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 +settags(n::IndexName, tags) = @set n.tags = to_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 +158,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 +189,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 +222,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) diff --git a/test/test_basics.jl b/test/test_basics.jl index cc9fcce4..d6bb6674 100644 --- a/test/test_basics.jl +++ b/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 From 0f7c7ce907e6ab502854937c3163163fd3577b38 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 15:26:04 -0400 Subject: [PATCH 3/4] Simplify tag input normalization settags no longer normalizes its argument: the IndexName tags field is typed, so `@set` enforces the type on construction, and its only callers (settag/unsettag) already build a SortedDict{Symbol,Symbol}. Split the conversion into to_symbol_pair (convert one pair's key and value to symbols) and to_tags, which accepts either a single bare Pair or a collection of pairs by dispatch, mirroring how Dict's constructor handles the same two shapes. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/index.jl | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/index.jl b/src/index.jl index 32d7b1c2..1b927d5f 100644 --- a/src/index.jl +++ b/src/index.jl @@ -29,13 +29,12 @@ function uniquename(rng::AbstractRNG, ::Type{<:IndexName}) return IndexName(rng) end -to_tags(tags::SortedDict{Symbol, Symbol}) = tags -to_tags(tag::Pair) = to_tags((tag,)) -function to_tags(tags) - return SortedDict{Symbol, Symbol}( - Symbol(first(p)) => Symbol(last(p)) for p in tags - ) -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) @@ -82,7 +81,7 @@ function Base.hash(n::IndexName, h::UInt) end setid(n::IndexName, id) = @set n.id = id -settags(n::IndexName, tags) = @set n.tags = to_tags(tags) +settags(n::IndexName, tags) = @set n.tags = tags setplev(n::IndexName, plev) = @set n.plev = plev hastag(n::IndexName, tagname) = haskey(tags_stored(n), Symbol(tagname)) From 06454a58d420e6a7fe35351f639ed373b742f326 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 15:53:13 -0400 Subject: [PATCH 4/4] Use findfirst for SortedDict lookups and credit prior art Replace the custom linear _searchsortedfirst, a misleading name for a linear scan, with findfirst, and drop the dead key/value converts. The keys stay sorted on insert, so lookups could later call searchsortedfirst for larger collections. Credit TensorKit.jl's SortedVectorDict as the model and note DataStructures.jl's SortedDict and Dictionaries.jl's ArrayDictionary as comparable designs. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/sorteddict.jl | 62 +++++++++++++++++++++-------------------------- 1 file changed, 28 insertions(+), 34 deletions(-) diff --git a/src/sorteddict.jl b/src/sorteddict.jl index 76673f3c..4a9bd1e9 100644 --- a/src/sorteddict.jl +++ b/src/sorteddict.jl @@ -1,3 +1,7 @@ +# 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} @@ -39,48 +43,38 @@ function Base.copy(d::SortedDict{K, V}) where {K, V} return SortedDict{K, V}(copy(keys(d)), copy(values(d))) end -# Linear scan: index of the first key not less than `key`. Fastest for small -# `length(d)`; swap in `searchsortedfirst` if a large-`n` use case appears. -function _searchsortedfirst(ks::Vector, key) - i = 1 - @inbounds while i <= length(ks) && isless(ks[i], key) - i += 1 - end - return i -end - -function Base.haskey(d::SortedDict{K}, key) where {K} - k = convert(K, key) - i = _searchsortedfirst(keys(d), k) - return i <= length(d) && isequal(keys(d)[i], k) +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{K}, key) where {K} - k = convert(K, key) - i = _searchsortedfirst(keys(d), k) - (i <= length(d) && isequal(keys(d)[i], k)) || throw(KeyError(key)) +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{K}, key, default) where {K} - k = convert(K, key) - i = _searchsortedfirst(keys(d), k) - return (i <= length(d) && isequal(keys(d)[i], k)) ? (@inbounds values(d)[i]) : default +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{K, V}, value, key) where {K, V} - k = convert(K, key) - v = convert(V, value) - i = _searchsortedfirst(keys(d), k) - if i <= length(d) && isequal(keys(d)[i], k) - @inbounds values(d)[i] = v +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 - insert!(keys(d), i, k) - insert!(values(d), i, v) + @inbounds values(d)[i] = value end return d end -function Base.delete!(d::SortedDict{K}, key) where {K} - k = convert(K, key) - i = _searchsortedfirst(keys(d), k) - if i <= length(d) && isequal(keys(d)[i], k) +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