diff --git a/src/ITensorBase.jl b/src/ITensorBase.jl index 278bfa1..74499ca 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/index.jl b/src/index.jl index 0d2dd77..1b927d5 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,23 +29,54 @@ 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 @@ -53,18 +84,22 @@ 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) diff --git a/src/sorteddict.jl b/src/sorteddict.jl new file mode 100644 index 0000000..4a9bd1e --- /dev/null +++ b/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 diff --git a/test/test_basics.jl b/test/test_basics.jl index cc9fcce..d6bb667 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 diff --git a/test/test_sorteddict.jl b/test/test_sorteddict.jl new file mode 100644 index 0000000..24df746 --- /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