Merge pull request #742 from AayushSabharwal/as/revert-weakset

feat: revert usage of `WeakCacheSet` for hashconsing

Author: AayushSabharwal

Project.toml

@@ -27,6 +27,7 @@ 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,6 +62,7 @@ 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

@@ -22,8 +22,9 @@ import TermInterface: iscall, isexpr, head, children,
 import ArrayInterface
 import ExproniconLite as EL
 import TaskLocalValues: TaskLocalValue
+using WeakValueDicts: WeakValueDict
 
-include("WeakCacheSets.jl")
+# include("WeakCacheSets.jl")
 
 include("cache.jl")
 Base.@deprecate istree iscall

src/types.jl

@@ -88,7 +88,7 @@ function exprtype(x::BasicSymbolic)
     end
 end
 
-const wcs = TaskLocalValue{WeakCacheSet{BasicSymbolic}}(WeakCacheSet{BasicSymbolic})
+const wvd = TaskLocalValue{WeakValueDict{UInt, BasicSymbolic}}(WeakValueDict{UInt, BasicSymbolic})
 
 # Same but different error messages
 @noinline error_on_type() = error("Internal error: unreachable reached!")
@@ -547,12 +547,21 @@ function BasicSymbolic(s::BasicSymbolic)::BasicSymbolic
     if !ENABLE_HASHCONSING[]
         return s
     end
-    cache = wcs[]
-    k = getkey!(cache, s)
-    if iszero(k.id[])
-        k.id[] = @atomic ID_COUNTER.x += 1
+
+    cache = wvd[]
+    h = hash2(s)
+    k = get!(cache, h, s)
+    if isequal_with_metadata(k, s)
+        if iszero(k.id[])
+            k.id[] = @atomic ID_COUNTER.x += 1
+        end
+        return k
+    else
+        if iszero(s.id[])
+            s.id[] = @atomic ID_COUNTER.x += 1
+        end
+        return s
     end
-    return k
 end
 
 function Sym{T}(name::Symbol; kw...) where {T}