Effects-tune PersitentDict (#52954)

To in particular allow elimination of dead PersistentDict constructions.

Author: Keno

base/array.jl

@@ -343,6 +343,9 @@ See also [`copy!`](@ref Base.copy!), [`copyto!`](@ref), [`deepcopy`](@ref).
 copy
 
 @eval function copy(a::Array{T}) where {T}
+    # `jl_genericmemory_copy_slice` only throws when the size exceeds the max allocation
+    # size, but since we're copying an existing array, we're guaranteed that this will not happen.
+    @_nothrow_meta
     ref = a.ref
     newmem = ccall(:jl_genericmemory_copy_slice, Ref{Memory{T}}, (Any, Ptr{Cvoid}, Int), ref.mem, ref.ptr_or_offset, length(a))
     return $(Expr(:new, :(typeof(a)), :(Core.memoryref(newmem)), :(a.size)))
@@ -1040,6 +1043,7 @@ end
 array_new_memory(mem::Memory, newlen::Int) = typeof(mem)(undef, newlen) # when implemented, this should attempt to first expand mem
 
 function _growbeg!(a::Vector, delta::Integer)
+    @_noub_meta
     delta = Int(delta)
     delta == 0 && return # avoid attempting to index off the end
     delta >= 0 || throw(ArgumentError("grow requires delta >= 0"))
@@ -1077,6 +1081,7 @@ function _growbeg!(a::Vector, delta::Integer)
 end
 
 function _growend!(a::Vector, delta::Integer)
+    @_noub_meta
     delta = Int(delta)
     delta >= 0 || throw(ArgumentError("grow requires delta >= 0"))
     ref = a.ref
@@ -1113,6 +1118,7 @@ function _growend!(a::Vector, delta::Integer)
 end
 
 function _growat!(a::Vector, i::Integer, delta::Integer)
+    @_terminates_globally_noub_meta
     delta = Int(delta)
     i = Int(i)
     i == 1 && return _growbeg!(a, delta)
@@ -1715,10 +1721,11 @@ julia> insert!(Any[1:6;], 3, "here")
 ```
 """
 function insert!(a::Array{T,1}, i::Integer, item) where T
+    @_noub_meta
     # Throw convert error before changing the shape of the array
     _item = item isa T ? item : convert(T, item)::T
     _growat!(a, i, 1)
-    # _growat! already did bound check
+    # :noub, because _growat! already did bound check
     @inbounds a[i] = _item
     return a
 end

base/compiler/tfuncs.jl

@@ -3127,6 +3127,8 @@ function foreigncall_effects(@specialize(abstract_eval), e::Expr)
     if name === :jl_alloc_genericmemory
         nothrow = new_genericmemory_nothrow(abstract_eval, args)
         return Effects(EFFECTS_TOTAL; consistent=CONSISTENT_IF_NOTRETURNED, nothrow)
+    elseif name === :jl_genericmemory_copy_slice
+        return Effects(EFFECTS_TOTAL; consistent=CONSISTENT_IF_NOTRETURNED, nothrow=false)
     end
     return EFFECTS_UNKNOWN
 end

base/dict.jl

@@ -901,14 +901,31 @@ struct PersistentDict{K,V} <: AbstractDict{K,V}
     @noinline function KeyValue.set(::Type{PersistentDict{K, V}}, ::Nothing, key, val) where {K, V}
         new{K, V}(HAMT.HAMT{K, V}(key => val))
     end
-    @noinline function KeyValue.set(dict::PersistentDict{K, V}, key, val) where {K, V}
+    @noinline @Base.assume_effects :effect_free function KeyValue.set(dict::PersistentDict{K, V}, key, val) where {K, V}
         trie = dict.trie
         h = HAMT.HashState(key)
         found, present, trie, i, bi, top, hs = HAMT.path(trie, key, h, #=persistent=# true)
         HAMT.insert!(found, present, trie, i, bi, hs, val)
         return new{K, V}(top)
     end
-    @noinline function KeyValue.set(dict::PersistentDict{K, V}, key) where {K, V}
+    @noinline @Base.assume_effects :nothrow :effect_free function KeyValue.set(dict::PersistentDict{K, V}, key::K, val::V) where {K, V}
+        trie = dict.trie
+        h = HAMT.HashState(key)
+        found, present, trie, i, bi, top, hs = HAMT.path(trie, key, h, #=persistent=# true)
+        HAMT.insert!(found, present, trie, i, bi, hs, val)
+        return new{K, V}(top)
+    end
+    @noinline @Base.assume_effects :effect_free function KeyValue.set(dict::PersistentDict{K, V}, key) where {K, V}
+        trie = dict.trie
+        h = HAMT.HashState(key)
+        found, present, trie, i, bi, top, _ = HAMT.path(trie, key, h, #=persistent=# true)
+        if found && present
+            deleteat!(trie.data, i)
+            HAMT.unset!(trie, bi)
+        end
+        return new{K, V}(top)
+    end
+    @noinline @Base.assume_effects :nothrow :effect_free function KeyValue.set(dict::PersistentDict{K, V}, key::K) where {K, V}
         trie = dict.trie
         h = HAMT.HashState(key)
         found, present, trie, i, bi, top, _ = HAMT.path(trie, key, h, #=persistent=# true)

base/essentials.jl

@@ -253,6 +253,32 @@ macro _terminates_globally_meta()
         #=:noub=#false,
         #=:noub_if_noinbounds=#false))
 end
+# can be used in place of `@assume_effects :terminates_globally :notaskstate` (supposed to be used for bootstrapping)
+macro _terminates_globally_notaskstate_meta()
+    return _is_internal(__module__) && Expr(:meta, Expr(:purity,
+        #=:consistent=#false,
+        #=:effect_free=#false,
+        #=:nothrow=#false,
+        #=:terminates_globally=#true,
+        #=:terminates_locally=#true,
+        #=:notaskstate=#true,
+        #=:inaccessiblememonly=#false,
+        #=:noub=#false,
+        #=:noub_if_noinbounds=#false))
+end
+# can be used in place of `@assume_effects :terminates_globally :noub` (supposed to be used for bootstrapping)
+macro _terminates_globally_noub_meta()
+    return _is_internal(__module__) && Expr(:meta, Expr(:purity,
+        #=:consistent=#false,
+        #=:effect_free=#false,
+        #=:nothrow=#false,
+        #=:terminates_globally=#true,
+        #=:terminates_locally=#true,
+        #=:notaskstate=#false,
+        #=:inaccessiblememonly=#false,
+        #=:noub=#true,
+        #=:noub_if_noinbounds=#false))
+end
 # can be used in place of `@assume_effects :effect_free :terminates_locally` (supposed to be used for bootstrapping)
 macro _effect_free_terminates_locally_meta()
     return _is_internal(__module__) && Expr(:meta, Expr(:purity,
@@ -279,6 +305,45 @@ macro _nothrow_noub_meta()
         #=:noub=#true,
         #=:noub_if_noinbounds=#false))
 end
+# can be used in place of `@assume_effects :nothrow` (supposed to be used for bootstrapping)
+macro _nothrow_meta()
+    return _is_internal(__module__) && Expr(:meta, Expr(:purity,
+        #=:consistent=#false,
+        #=:effect_free=#false,
+        #=:nothrow=#true,
+        #=:terminates_globally=#false,
+        #=:terminates_locally=#false,
+        #=:notaskstate=#false,
+        #=:inaccessiblememonly=#false,
+        #=:noub=#false,
+        #=:noub_if_noinbounds=#false))
+end
+# can be used in place of `@assume_effects :nothrow` (supposed to be used for bootstrapping)
+macro _noub_meta()
+    return _is_internal(__module__) && Expr(:meta, Expr(:purity,
+        #=:consistent=#false,
+        #=:effect_free=#false,
+        #=:nothrow=#false,
+        #=:terminates_globally=#false,
+        #=:terminates_locally=#false,
+        #=:notaskstate=#false,
+        #=:inaccessiblememonly=#false,
+        #=:noub=#true,
+        #=:noub_if_noinbounds=#false))
+end
+# can be used in place of `@assume_effects :notaskstate` (supposed to be used for bootstrapping)
+macro _notaskstate_meta()
+    return _is_internal(__module__) && Expr(:meta, Expr(:purity,
+        #=:consistent=#false,
+        #=:effect_free=#false,
+        #=:nothrow=#false,
+        #=:terminates_globally=#false,
+        #=:terminates_locally=#false,
+        #=:notaskstate=#true,
+        #=:inaccessiblememonly=#false,
+        #=:noub=#false,
+        #=:noub_if_noinbounds=#false))
+end
 # can be used in place of `@assume_effects :noub_if_noinbounds` (supposed to be used for bootstrapping)
 macro _noub_if_noinbounds_meta()
     return _is_internal(__module__) && Expr(:meta, Expr(:purity,

base/genericmemory.jl

@@ -70,7 +70,7 @@ isassigned(a::GenericMemoryRef) = memoryref_isassigned(a, :not_atomic, @_boundsc
 
 ## copy ##
 function unsafe_copyto!(dest::MemoryRef{T}, src::MemoryRef{T}, n) where {T}
-    @_terminates_globally_meta
+    @_terminates_globally_notaskstate_meta
     n == 0 && return dest
     @boundscheck GenericMemoryRef(dest, n), GenericMemoryRef(src, n)
     ccall(:jl_genericmemory_copyto, Cvoid, (Any, Ptr{Cvoid}, Any, Ptr{Cvoid}, Int), dest.mem, dest.ptr_or_offset, src.mem, src.ptr_or_offset, Int(n))

base/hamt.jl

@@ -92,7 +92,7 @@ struct HashState{K}
 end
 HashState(key) = HashState(key, objectid(key), 0, 0)
 # Reconstruct
-function HashState(other::HashState, key)
+Base.@assume_effects :terminates_locally function HashState(other::HashState, key)
     h = HashState(key)
     while h.depth !== other.depth
         h = next(h)
@@ -103,7 +103,8 @@ end
 function next(h::HashState)
     depth = h.depth + 1
     shift = h.shift + BITS_PER_LEVEL
-    @assert h.shift <= MAX_SHIFT
+    # Assert disabled for effect precision
+    # @assert h.shift <= MAX_SHIFT
     if shift > MAX_SHIFT
         # Note we use `UInt(depth ÷ BITS_PER_LEVEL)` to seed the hash function
         # the hash docs, do we need to hash `UInt(depth ÷ BITS_PER_LEVEL)` first?
@@ -154,7 +155,7 @@ as the current `level`.
 If a copy function is provided `copyf` use the return `top` for the
 new persistent tree.
 """
-@inline function path(trie::HAMT{K,V}, key, h::HashState, copy=false) where {K, V}
+@inline @Base.assume_effects :noub :terminates_locally function path(trie::HAMT{K,V}, key, h::HashState, copy=false) where {K, V}
     if copy
         trie = top = HAMT{K,V}(Base.copy(trie.data), trie.bitmap)
     else
@@ -172,6 +173,7 @@ new persistent tree.
             end
             if copy
                 next = HAMT{K,V}(Base.copy(next.data), next.bitmap)
+                # :noub because entry_index is guaranteed to be inbounds for trie.data
                 @inbounds trie.data[i] = next
             end
             trie = next::HAMT{K,V}
@@ -187,7 +189,7 @@ end
 Internal function that given an obtained path, either set the value
 or grows the HAMT by inserting a new trie instead.
 """
-@inline function insert!(found, present, trie::HAMT{K,V}, i, bi, h, val) where {K,V}
+@inline @Base.assume_effects :terminates_locally function insert!(found, present, trie::HAMT{K,V}, i, bi, h, val) where {K,V}
     if found # we found a slot, just set it to the new leaf
         # replace or insert
         if present # replace

test/compiler/irpasses.jl

@@ -1561,7 +1561,10 @@ function persistent_dict_elim_multiple()
     return b[:a]
 end
 @test_broken fully_eliminated(persistent_dict_elim_multiple)
-@test code_typed(persistent_dict_elim_multiple)[1][1].code[end] == Core.ReturnNode(1)
+let code = code_typed(persistent_dict_elim_multiple)[1][1].code
+    @test count(x->isexpr(x, :invoke), code) == 0
+    @test code[end] == Core.ReturnNode(1)
+end
 
 function persistent_dict_elim_multiple_phi(c::Bool)
     if c