Merge pull request #713 from JuliaSymbolics/b/WeakKeyDict

Replace `WeakValueDict` with `WeakKeyDict` for hash consing

Author: ChrisRackauckas

Project.toml

@@ -27,7 +27,6 @@ TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34"
 TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Unityper = "a7c27f48-0311-42f6-a7f8-2c11e75eb415"
-WeakValueDicts = "897b6980-f191-5a31-bcb0-bf3c4585e0c1"
 
 [weakdeps]
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
@@ -61,7 +60,6 @@ TaskLocalValues = "0.1.2"
 TermInterface = "2.0"
 TimerOutputs = "0.5"
 Unityper = "0.1.2"
-WeakValueDicts = "0.1.0"
 julia = "1.10"
 
 [extras]

src/SymbolicUtils.jl

@@ -20,7 +20,6 @@ import TermInterface: iscall, isexpr, head, children,
                       operation, arguments, metadata, maketerm, sorted_arguments
 # For ReverseDiffExt
 import ArrayInterface
-using WeakValueDicts: WeakValueDict
 import ExproniconLite as EL
 import TaskLocalValues: TaskLocalValue
 

src/types.jl

@@ -23,39 +23,39 @@ const EMPTY_DICT_T = typeof(EMPTY_DICT)
 const ENABLE_HASHCONSING = Ref(true)
 
 @compactify show_methods=false begin
-    @abstract mutable struct BasicSymbolic{T} <: Symbolic{T}
+    @abstract struct BasicSymbolic{T} <: Symbolic{T}
         metadata::Metadata     = NO_METADATA
     end
-    mutable struct Sym{T} <: BasicSymbolic{T}
+    struct Sym{T} <: BasicSymbolic{T}
         name::Symbol           = :OOF
     end
-    mutable struct Term{T} <: BasicSymbolic{T}
+    struct Term{T} <: BasicSymbolic{T}
         f::Any                 = identity  # base/num if Pow; issorted if Add/Dict
         arguments::Vector{Any} = EMPTY_ARGS
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
     end
-    mutable struct Mul{T} <: BasicSymbolic{T}
+    struct Mul{T} <: BasicSymbolic{T}
         coeff::Any             = 0         # exp/den if Pow
         dict::EMPTY_DICT_T     = EMPTY_DICT
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
         arguments::Vector{Any} = EMPTY_ARGS
     end
-    mutable struct Add{T} <: BasicSymbolic{T}
+    struct Add{T} <: BasicSymbolic{T}
         coeff::Any             = 0         # exp/den if Pow
         dict::EMPTY_DICT_T     = EMPTY_DICT
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
         arguments::Vector{Any} = EMPTY_ARGS
     end
-    mutable struct Div{T} <: BasicSymbolic{T}
+    struct Div{T} <: BasicSymbolic{T}
         num::Any               = 1
         den::Any               = 1
         simplified::Bool       = false
         arguments::Vector{Any} = EMPTY_ARGS
     end
-    mutable struct Pow{T} <: BasicSymbolic{T}
+    struct Pow{T} <: BasicSymbolic{T}
         base::Any              = 1
         exp::Any               = 1
         arguments::Vector{Any} = EMPTY_ARGS
@@ -86,7 +86,15 @@ function exprtype(x::BasicSymbolic)
     end
 end
 
-const wvd = TaskLocalValue{WeakValueDict{UInt, BasicSymbolic}}(WeakValueDict{UInt, BasicSymbolic})
+mutable struct HashConsingWrapper
+    bs::BasicSymbolic
+end
+
+Base.hash(x::HashConsingWrapper, h::UInt) = hash2(x.bs, h)
+
+Base.isequal(x::HashConsingWrapper, y::HashConsingWrapper) = isequal_with_metadata(x.bs, y.bs)
+
+const wkd = TaskLocalValue{WeakKeyDict{HashConsingWrapper, Nothing}}(WeakKeyDict{HashConsingWrapper, Nothing})
 
 # Same but different error messages
 @noinline error_on_type() = error("Internal error: unreachable reached!")
@@ -522,15 +530,15 @@ Implements hash consing (flyweight design pattern) for `BasicSymbolic` objects.
 
 This function checks if an equivalent `BasicSymbolic` object already exists. It uses a 
 custom hash function (`hash2`) incorporating metadata and symtypes to search for existing 
-objects in a `WeakValueDict` (`wvd`).  Due to the possibility of hash collisions (where 
+objects in a `WeakKeyDict` (`wkd`). Due to the possibility of hash collisions (where 
 different objects produce the same hash), a custom equality check (`isequal_with_metadata`) 
 which includes metadata comparison, is used to confirm the equivalence of objects with 
 matching hashes. If an equivalent object is found, the existing object is returned; 
 otherwise, the input `s` is returned. This reduces memory usage, improves compilation time 
 for runtime code generation, and supports built-in common subexpression elimination, 
 particularly when working with symbolic objects with metadata.
 
-Using a `WeakValueDict` ensures that only weak references to `BasicSymbolic` objects are 
+Using a `WeakKeyDict` ensures that only weak references to `BasicSymbolic` objects are 
 stored, allowing objects that are no longer strongly referenced to be garbage collected. 
 Custom functions `hash2` and `isequal_with_metadata` are used instead of `Base.hash` and 
 `Base.isequal` to accommodate metadata without disrupting existing tests reliant on the 
@@ -540,12 +548,14 @@ function BasicSymbolic(s::BasicSymbolic)::BasicSymbolic
     if !ENABLE_HASHCONSING[]
         return s
     end
-    h = hash2(s)
-    t = get!(wvd[], h, s)
-    if t === s || isequal_with_metadata(t, s)
-        t
+    cache = wkd[]
+    hcw = HashConsingWrapper(s)
+    k = getkey(cache, hcw, nothing)
+    if isnothing(k)
+        cache[hcw] = nothing
+        return s
     else
-        s
+        return k.bs
     end
 end
 

test/hash_consing.jl

@@ -133,9 +133,9 @@ end
 
 @testset "Hashconsing can be toggled" begin
     SymbolicUtils.ENABLE_HASHCONSING[] = false
-    name = gensym(:x)
-    x1 = only(@eval @syms $name)
-    x2 = only(@eval @syms $name)
+    @syms a b
+    x1 = a + b
+    x2 = a + b
     @test x1 !== x2
     SymbolicUtils.ENABLE_HASHCONSING[] = true
 end