fix: use a GC allocation for hashconsed symbolic ID

Author: AayushSabharwal

src/cache.jl

@@ -5,6 +5,8 @@ Sentinel value used for a cache miss, since cached functions may return `nothing
 """
 struct CacheSentinel end
 
+mutable struct IDType end
+
 """
     $(TYPEDEF)
 
@@ -13,7 +15,7 @@ Struct wrapping the `id` of a `BasicSymbolic`, since arguments annotated
 up a symbolic or a `UInt`.
 """
 struct SymbolicKey
-    id::UInt64
+    id::IDType
 end
 
 """
@@ -23,7 +25,17 @@ 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
+function get_cache_key(x)
+    if x isa BasicSymbolic
+        id = x.id[]
+        if id === nothing
+            return CacheSentinel()
+        end
+        return SymbolicKey(id)
+    else
+        x
+    end
+end
 
 """
     associated_cache(fn)
@@ -241,6 +253,7 @@ macro cache(args...)
     argexprs = []
     # The name of the variable storing the result of looking up the cache
     cache_value_name = :val
+    valid_key_condition = :(true)
     # The condition for a cache hit
     cache_hit_condition = :(!($cache_value_name isa $CacheSentinel))
 
@@ -254,6 +267,7 @@ macro cache(args...)
             push!(keyexprs, :($get_cache_key($arg)))
             push!(argexprs, arg)
             push!(keytypes, Any)
+            valid_key_condition = :($valid_key_condition && !(key[$(length(keyexprs))] isa $CacheSentinel))
             continue
         end
         argname, Texpr = arg.args
@@ -263,6 +277,7 @@ macro cache(args...)
             # if the type is `Any`, branch on it being a `BasicSymbolic`
             push!(keyexprs, :($get_cache_key($argname)))
             push!(keytypes, Any)
+            valid_key_condition = :($valid_key_condition && !(key[$(length(keyexprs))] isa $CacheSentinel))
             continue
         end
 
@@ -275,6 +290,7 @@ macro cache(args...)
             push!(keytypes, Union{keyTs...})
             if maybe_basicsymbolic
                 push!(keyexprs, :($get_cache_key($argname)))
+                valid_key_condition = :($valid_key_condition && !(key[$(length(keyexprs))] isa $CacheSentinel))
             else
                 push!(keyexprs, argname)
             end
@@ -286,6 +302,7 @@ macro cache(args...)
         if T <: BasicSymbolic
             push!(keytypes, SymbolicKey) 
             push!(keyexprs, :($get_cache_key($argname)))
+            valid_key_condition = :($valid_key_condition && !(key[$(length(keyexprs))] isa $CacheSentinel))
         else
             push!(keytypes, T)
             push!(keyexprs, argname)
@@ -345,6 +362,9 @@ macro cache(args...)
         if $conditions
             # construct the `Tuple` key
             key = $keyexpr
+            if !($valid_key_condition)
+                return $innercall
+            end
             # get the cache and stats from the `TaskLocalValue` to avoid accessing
             # `task_local_storage` repeatedly.
             cachedict, cachestats = $cachename.tlv[]

src/types.jl

@@ -22,11 +22,13 @@ const EMPTY_HASH = RefValue(UInt(0))
 const EMPTY_DICT = sdict()
 const EMPTY_DICT_T = typeof(EMPTY_DICT)
 const ENABLE_HASHCONSING = Ref(true)
+const TID = Union{IDType, Nothing}
+const DID = nothing
 
 @compactify show_methods=false begin
     @abstract mutable struct BasicSymbolic{T} <: Symbolic{T}
         metadata::Metadata     = NO_METADATA
-        id::RefValue{UInt64} = Ref{UInt64}(0)
+        id::RefValue{TID} = Ref{TID}(DID)
     end
     mutable struct Sym{T} <: BasicSymbolic{T}
         name::Symbol           = :OOF
@@ -88,8 +90,6 @@ function exprtype(x::BasicSymbolic)
     end
 end
 
-const wvd = TaskLocalValue{WeakValueDict{UInt, BasicSymbolic}}(WeakValueDict{UInt, BasicSymbolic})
-
 # Same but different error messages
 @noinline error_on_type() = error("Internal error: unreachable reached!")
 @noinline error_sym() = error("Sym doesn't have a operation or arguments!")
@@ -108,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)), 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))
+        Term => Term{T}(nt_new.f, nt_new.arguments; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{TID}(DID))
+        Add => Add(T, nt_new.coeff, nt_new.dict; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{TID}(DID))
+        Mul => Mul(T, nt_new.coeff, nt_new.dict; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{TID}(DID))
+        Div => Div{T}(nt_new.num, nt_new.den, nt_new.simplified; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{TID}(DID))
+        Pow => Pow{T}(nt_new.base, nt_new.exp; nt_new..., hash = RefValue(UInt(0)), hash2 = RefValue(UInt(0)), id = Ref{TID}(DID))
         _ => Unityper.rt_constructor(obj){T}(;nt_new...)
     end
 end
@@ -516,11 +516,11 @@ end
 ### Constructors
 ###
 
-mutable struct AtomicIDCounter
-    @atomic x::UInt64
-end
+const wvd = TaskLocalValue{WeakValueDict{UInt, BasicSymbolic}}(WeakValueDict{UInt, BasicSymbolic})
 
-const ID_COUNTER = AtomicIDCounter(0)
+function generate_id()
+    return IDType()
+end
 
 """
 $(TYPEDSIGNATURES)
@@ -552,27 +552,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
+        if isnothing(k.id[])
+            k.id[] = generate_id()
         end
         return k
     else
-        if iszero(s.id[])
-            s.id[] = @atomic ID_COUNTER.x += 1
+        if isnothing(s.id[])
+            s.id[] = generate_id()
         end
         return s
     end
 end
 
 function Sym{T}(name::Symbol; kw...) where {T}
-    s = Sym{T}(; name, kw..., id = Ref{UInt}(0))
+    s = Sym{T}(; name, kw..., id = Ref{TID}(DID))
     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..., id = Ref{UInt64}(0))
+    s = Term{T}(;f=f, arguments=args, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), kw..., id = Ref{TID}(DID))
     BasicSymbolic(s)
 end
 
@@ -602,7 +602,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..., id = Ref{UInt64}(0))
+    s = Add{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw..., id = Ref{TID}(DID))
     BasicSymbolic(s)
 end
 
@@ -620,7 +620,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..., id = Ref{UInt64}(0))
+        s = Mul{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=SmallV{Any}(), kw..., id = Ref{TID}(DID))
         BasicSymbolic(s)
     end
 end
@@ -688,7 +688,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, id = Ref{UInt64}(0))
+    s = Div{T}(; num=n, den=d, simplified, arguments=SmallV{Any}(), metadata, id = Ref{TID}(DID))
     BasicSymbolic(s)
 end
 
@@ -708,7 +708,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, id = Ref{UInt64}(0))
+    s = Pow{T}(; base=a, exp=b, arguments=SmallV{Any}(), metadata, id = Ref{TID}(DID))
     BasicSymbolic(s)
 end