feat: use a global atomic UInt64 to give each BasicSymbolic a unique ID

Author: AayushSabharwal

src/cache.jl

@@ -8,14 +8,23 @@ struct CacheSentinel end
 """
     $(TYPEDEF)
 
-Struct wrapping the `objectid` of a `BasicSymbolic`, since arguments annotated
+Struct wrapping the `id` of a `BasicSymbolic`, since arguments annotated
 `::Union{BasicSymbolic, UInt}` would not be able to differentiate between looking
 up a symbolic or a `UInt`.
 """
 struct SymbolicKey
-    id::UInt
+    id::UInt64
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+The key stored in the cache for a particular value. Returns a `SymbolicKey` for
+`BasicSymbolic` and is the identity function otherwise.
+"""
+# can't dispatch because `BasicSymbolic` isn't defined here
+get_cache_key(x) = x isa BasicSymbolic ? SymbolicKey(x.id[]) : x
+
 """
     associated_cache(fn)
 
@@ -234,10 +243,6 @@ macro cache(args...)
     cache_value_name = :val
     # The condition for a cache hit
     cache_hit_condition = :(!($cache_value_name isa $CacheSentinel))
-    # Type of additional data stored with cached result. Used to compare
-    # equality of `BasicSymbolic` arguments, since `objectid` is a hash.
-    cache_additional_types = []
-    cache_additional_values = []
 
     for arg in fn.args
         # handle arguments with defaults
@@ -246,24 +251,18 @@ macro cache(args...)
         end
         if !Meta.isexpr(arg, :(::))
             # if the type is `Any`, branch on it being a `BasicSymbolic`
-            push!(keyexprs, :($arg isa BasicSymbolic ? $SymbolicKey(objectid($arg)) : $arg))
+            push!(keyexprs, :($get_cache_key($arg)))
             push!(argexprs, arg)
             push!(keytypes, Any)
-            push!(cache_additional_types, Any)
-            push!(cache_additional_values, arg)
-            cache_hit_condition = :($cache_hit_condition && (!($arg isa BasicSymbolic) || $arg === $cache_value_name[$(length(cache_additional_values))]))
             continue
         end
         argname, Texpr = arg.args
         push!(argexprs, argname)
 
         if Texpr == :Any
             # if the type is `Any`, branch on it being a `BasicSymbolic`
-            push!(keyexprs, :($argname isa BasicSymbolic ? $SymbolicKey(objectid($argname)) : $argname))
+            push!(keyexprs, :($get_cache_key($argname)))
             push!(keytypes, Any)
-            push!(cache_additional_types, Any)
-            push!(cache_additional_values, argname)
-            cache_hit_condition = :($cache_hit_condition && (!($argname isa BasicSymbolic) || $argname === $cache_value_name[$(length(cache_additional_values))]))
             continue
         end
 
@@ -275,10 +274,7 @@ macro cache(args...)
             maybe_basicsymbolic = any(x -> x <: BasicSymbolic, Ts)
             push!(keytypes, Union{keyTs...})
             if maybe_basicsymbolic
-                push!(keyexprs, :($argname isa BasicSymbolic ? $SymbolicKey(objectid($argname)) : $argname))
-                push!(cache_additional_types, Texpr)
-                push!(cache_additional_values, argname)
-                cache_hit_condition = :($cache_hit_condition && (!($argname isa BasicSymbolic) || $argname === $cache_value_name[$(length(cache_additional_values))]))
+                push!(keyexprs, :($get_cache_key($argname)))
             else
                 push!(keyexprs, argname)
             end
@@ -289,10 +285,7 @@ macro cache(args...)
         T = Base.eval(__module__, Texpr)
         if T <: BasicSymbolic
             push!(keytypes, SymbolicKey) 
-            push!(keyexprs, :($SymbolicKey(objectid($argname))))
-            push!(cache_additional_types, T)
-            push!(cache_additional_values, argname)
-            cache_hit_condition = :($cache_hit_condition && $argname === $cache_value_name[$(length(cache_additional_values))])
+            push!(keyexprs, :($get_cache_key($argname)))
         else
             push!(keytypes, T)
             push!(keyexprs, argname)
@@ -313,9 +306,7 @@ macro cache(args...)
     # construct an expression for the type of the cache keys
     keyT = Expr(:curly, Tuple)
     append!(keyT.args, keytypes)
-    valT = Expr(:curly, Tuple)
-    append!(valT.args, cache_additional_types)
-    push!(valT.args, rettype)
+    valT = rettype
     # the type of the cache
     cacheT = :(Dict{$keyT, $valT})
     # type of the `TaskLocalValue`
@@ -364,7 +355,7 @@ macro cache(args...)
             if $cache_hit_condition
                 # cache hit
                 cachestats.hits += 1
-                return $cache_value_name[end]
+                return $cache_value_name
             end
             # cache miss
             cachestats.misses += 1
@@ -375,8 +366,8 @@ macro cache(args...)
                 $(filter!)($cachename, cachedict)
             end
             # add to cache
-            cachedict[key] = ($(cache_additional_values...), val)
-            return val
+            cachedict[key] = $cache_value_name
+            return $cache_value_name
         end
 
         # if we're not doing caching

src/types.jl

@@ -26,6 +26,7 @@ const ENABLE_HASHCONSING = Ref(true)
 @compactify show_methods=false begin
     @abstract mutable struct BasicSymbolic{T} <: Symbolic{T}
         metadata::Metadata     = NO_METADATA
+        id::RefValue{UInt64} = Ref{UInt64}(0)
     end
     mutable struct Sym{T} <: BasicSymbolic{T}
         name::Symbol           = :OOF
@@ -107,11 +108,11 @@ function ConstructionBase.setproperties(obj::BasicSymbolic{T}, patch::NamedTuple
     # Call outer constructor because hash consing cannot be applied in inner constructor
     @compactified obj::BasicSymbolic begin
         Sym => Sym{T}(nt_new.name; nt_new...)
-        Term => Term{T}(nt_new.f, nt_new.arguments; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)))
-        Add => Add(T, nt_new.coeff, nt_new.dict; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)))
-        Mul => Mul(T, nt_new.coeff, nt_new.dict; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)))
-        Div => Div{T}(nt_new.num, nt_new.den, nt_new.simplified; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)))
-        Pow => Pow{T}(nt_new.base, nt_new.exp; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)))
+        Term => Term{T}(nt_new.f, nt_new.arguments; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{UInt64}(0))
+        Add => Add(T, nt_new.coeff, nt_new.dict; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{UInt64}(0))
+        Mul => Mul(T, nt_new.coeff, nt_new.dict; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{UInt64}(0))
+        Div => Div{T}(nt_new.num, nt_new.den, nt_new.simplified; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{UInt64}(0))
+        Pow => Pow{T}(nt_new.base, nt_new.exp; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{UInt64}(0))
         _ => Unityper.rt_constructor(obj){T}(;nt_new...)
     end
 end
@@ -255,6 +256,7 @@ end
 
 function Base.isequal(a::BasicSymbolic{T}, b::BasicSymbolic{S}) where {T,S}
     a === b && return true
+    a.id == b.id && a.id != 0 && return true
 
     E = exprtype(a)
     E === exprtype(b) || return false
@@ -298,6 +300,7 @@ function.
 """
 function isequal_with_metadata(a::BasicSymbolic{T}, b::BasicSymbolic{S})::Bool where {T, S}
     a === b && return true
+    a.id == b.id && a.id != 0 && return true
 
     E = exprtype(a)
     E === exprtype(b) || return false
@@ -513,6 +516,12 @@ end
 ### Constructors
 ###
 
+mutable struct AtomicIDCounter
+    @atomic x::UInt64
+end
+
+const ID_COUNTER = AtomicIDCounter(0)
+
 """
 $(TYPEDSIGNATURES)
 
@@ -542,21 +551,27 @@ function BasicSymbolic(s::BasicSymbolic)::BasicSymbolic
     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
 end
 
 function Sym{T}(name::Symbol; kw...) where {T}
-    s = Sym{T}(; name, kw...)
+    s = Sym{T}(; name, kw..., id = Ref{UInt}(0))
     BasicSymbolic(s)
 end
 
 function Term{T}(f, args; kw...) where T
     args = SmallV{Any}(args)
 
-    s = Term{T}(;f=f, arguments=args, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), kw...)
+    s = Term{T}(;f=f, arguments=args, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), kw..., id = Ref{UInt64}(0))
     BasicSymbolic(s)
 end
 
@@ -586,7 +601,7 @@ function Add(::Type{T}, coeff, dict; metadata=NO_METADATA, kw...) where T
         end
     end
 
-    s = Add{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw...)
+    s = Add{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw..., id = Ref{UInt64}(0))
     BasicSymbolic(s)
 end
 
@@ -604,7 +619,7 @@ function Mul(T, a, b; metadata=NO_METADATA, kw...)
     else
         coeff = a
         dict = b
-        s = Mul{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw...)
+        s = Mul{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw..., id = Ref{UInt64}(0))
         BasicSymbolic(s)
     end
 end
@@ -672,7 +687,7 @@ function Div{T}(n, d, simplified=false; metadata=nothing, kwargs...) where {T}
         end
     end
 
-    s = Div{T}(; num=n, den=d, simplified, arguments=SmallV{Any}(), metadata)
+    s = Div{T}(; num=n, den=d, simplified, arguments=SmallV{Any}(), metadata, id = Ref{UInt64}(0))
     BasicSymbolic(s)
 end
 
@@ -692,7 +707,7 @@ function Pow{T}(a, b; metadata=NO_METADATA, kwargs...) where {T}
     b = unwrap(b)
     _iszero(b) && return 1
     _isone(b) && return a
-    s = Pow{T}(; base=a, exp=b, arguments=SmallV{Any}(), metadata)
+    s = Pow{T}(; base=a, exp=b, arguments=SmallV{Any}(), metadata, id = Ref{UInt64}(0))
     BasicSymbolic(s)
 end
 

test/cache_macro.jl

@@ -1,6 +1,6 @@
 using SymbolicUtils
 using SymbolicUtils: BasicSymbolic, @cache, associated_cache, set_limit!, get_limit,
-                     clear_cache!, SymbolicKey, metadata, maketerm
+                     clear_cache!, SymbolicKey, metadata, maketerm, get_cache_key
 using OhMyThreads: tmap
 using Random
 
@@ -14,9 +14,9 @@ end
     @test isequal(val, 2x + 1)
     cachestruct = associated_cache(f1)
     cache, stats = cachestruct.tlv[]
-    @test cache isa Dict{Tuple{SymbolicKey}, Tuple{BasicSymbolic, BasicSymbolic}}
+    @test cache isa Dict{Tuple{SymbolicKey}, BasicSymbolic}
     @test length(cache) == 1
-    @test cache[(SymbolicKey(objectid(x)),)][end] === val
+    @test cache[(get_cache_key(x),)] === val
     @test stats.hits == 0
     @test stats.misses == 1
     f1(x)
@@ -76,20 +76,20 @@ end
     @test isequal(val, 2x + 1)
     cachestruct = associated_cache(f2)
     cache, stats = cachestruct.tlv[]
-    @test cache isa Dict{Tuple{Union{SymbolicKey, UInt}}, NTuple{2, Union{BasicSymbolic, UInt}}}
+    @test cache isa Dict{Tuple{Union{SymbolicKey, UInt}}, Union{BasicSymbolic, UInt}}
     @test length(cache) == 1
-    @test cache[(SymbolicKey(objectid(x)),)][end] === val
+    @test cache[(get_cache_key(x),)] === val
     @test stats.hits == 0
     @test stats.misses == 1
     f2(x)
     @test stats.hits == 1
     @test stats.misses == 1
 
-    y = objectid(x)
+    y = get_cache_key(x).id
     val = f2(y)
     @test val == 2y + 1
     @test length(cache) == 2
-    @test cache[(y,)][end] == val
+    @test cache[(y,)] == val
     @test stats.misses == 2
 
     clear_cache!(f2)
@@ -111,9 +111,9 @@ end
     @test isequal(val, 2x + 1)
     cachestruct = associated_cache(fn)
     cache, stats = cachestruct.tlv[]
-    @test cache isa Dict{Tuple{Any}, Tuple{Any, Union{BasicSymbolic, Int}}}
+    @test cache isa Dict{Tuple{Any}, Union{BasicSymbolic, Int}}
     @test length(cache) == 1
-    @test cache[(SymbolicKey(objectid(x)),)][end] === val
+    @test cache[(get_cache_key(x),)] === val
     @test stats.hits == 0
     @test stats.misses == 1
     fn(x)
@@ -161,42 +161,6 @@ end
     @test isequal(result, truevals)
 end
 
-@cache function f5(x::BasicSymbolic, y::Union{BasicSymbolic, Int}, z)::BasicSymbolic
-    return x + y + z
-end
-
-# temporary definition to induce objectid collisions
-Base.objectid(x::BasicSymbolic) = 0x42
-
-@testset "`objectid` collision handling" begin
-    @syms x y z
-    @test objectid(x) == objectid(y) == objectid(z) == 0x42
-    cachestruct = associated_cache(f5)
-    cache, stats = cachestruct.tlv[]
-    val = f5(x, 1, 2)
-    @test isequal(val, x + 3)
-    @test length(cache) == 1
-    @test stats.misses == 1
-    val2 = f5(y, 1, 2)
-    @test isequal(val2, y + 3)
-    @test length(cache) == 1
-    @test stats.misses == 2
-
-    clear_cache!(f5)
-    val = f5(x, y, z)
-    @test isequal(val, x + y + z)
-    @test length(cache) == 1
-    @test stats.misses == 1
-    val2 = f5(y, 2z, x)
-    @test isequal(val2, x + y + 2z)
-    @test length(cache) == 1
-    @test stats.misses == 2
-end
-
-Base.delete_method(only(methods(objectid, @__MODULE__)))
-@syms x
-@test objectid(x) != 0x42
-
 @cache limit = 10 retain_fraction = 0.1 function f6(x::BasicSymbolic, y::Union{BasicSymbolic, Int}, z)::BasicSymbolic
     return x + y + z
 end