Merge pull request #764 from AayushSabharwal/as/wcs

refactor: use new `WeakCacheSet` for hashconsing

Author: AayushSabharwal

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"
@@ -62,7 +61,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

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

src/WeakCacheSets.jl

@@ -1,3 +1,9 @@
+# Copied from WeakCacheSets.jl, modified to use hash2 and isequal_with_metadata instead of
+# the standard `hash` and `isequal`.
+module WeakCacheSets
+
+using SymbolicUtils: hash2, isequal_with_metadata
+
 # These can be changed, to trade off better performance for space
 const maxallowedprobe = 16
 const maxprobeshift   = 6
@@ -229,3 +235,7 @@ function Base.show(io::IO, t::WeakCacheSet{K}) where K
         end
     print(io, ')')
 end
+
+export WeakCacheSet, getkey!
+
+end # module WeakCacheSets

src/types.jl

@@ -516,7 +516,7 @@ end
 ### Constructors
 ###
 
-const wvd = TaskLocalValue{WeakValueDict{UInt, BasicSymbolic}}(WeakValueDict{UInt, BasicSymbolic})
+const wvd = TaskLocalValue{WeakCacheSet{BasicSymbolic}}(WeakCacheSet{BasicSymbolic})
 
 function generate_id()
     return IDType()
@@ -549,19 +549,11 @@ function BasicSymbolic(s::BasicSymbolic)::BasicSymbolic
     end
 
     cache = wvd[]
-    h = hash2(s)
-    k = get!(cache, h, s)
-    if isequal_with_metadata(k, s)
-        if isnothing(k.id[])
-            k.id[] = generate_id()
-        end
-        return k
-    else
-        if isnothing(s.id[])
-            s.id[] = generate_id()
-        end
-        return s
+    k = getkey!(cache, s)
+    if k.id[] === nothing
+        k.id[] = generate_id()
     end
+    return k
 end
 
 function Sym{T}(name::Symbol; kw...) where {T}