Merge pull request #766 from AayushSabharwal/as/fix-threading

fix: ensure equivalent expressions on different tasks are compared properly

Author: AayushSabharwal

src/cache.jl

@@ -27,7 +27,7 @@ The key stored in the cache for a particular value. Returns a `SymbolicKey` for
 # can't dispatch because `BasicSymbolic` isn't defined here
 function get_cache_key(x)
     if x isa BasicSymbolic
-        id = x.id
+        id = x.id[2]
         if id === nothing
             return CacheSentinel()
         end

src/types.jl

@@ -35,7 +35,7 @@ const ROArgsT = ReadOnlyVector{Any, ArgsT}
 const ACDict{K, V} = Dict{K, V}
 const ShapeVecT = SmallV{UnitRange{Int}}
 const ShapeT = Union{Unknown, ShapeVecT}
-const IdentT = Union{IDType, Nothing}
+const IdentT = Union{Tuple{UInt, IDType}, Tuple{Nothing, Nothing}}
 
 """
     Enum used to differentiate between variants of `BasicSymbolicImpl.ACTerm`.
@@ -171,22 +171,22 @@ override_properties(obj::BSImpl.Type) = override_properties(MData.variant_type(o
 
 function override_properties(obj::Type{<:BSImpl.Variant})
     @match obj begin
-        ::Type{<:BSImpl.Sym} => (; id = nothing, hash2 = 0)
-        ::Type{<:BSImpl.Term} => (; id = nothing, hash = 0, hash2 = 0)
-        ::Type{<:BSImpl.AddOrMul} => (; id = nothing, hash = 0, hash2 = 0)
-        ::Type{<:BSImpl.Div} => (; id = nothing, hash2 = 0)
-        ::Type{<:BSImpl.Pow} => (; id = nothing, hash2 = 0)
+        ::Type{<:BSImpl.Sym} => (; id = (nothing, nothing), hash2 = 0)
+        ::Type{<:BSImpl.Term} => (; id = (nothing, nothing), hash = 0, hash2 = 0)
+        ::Type{<:BSImpl.AddOrMul} => (; id = (nothing, nothing), hash = 0, hash2 = 0)
+        ::Type{<:BSImpl.Div} => (; id = (nothing, nothing), hash2 = 0)
+        ::Type{<:BSImpl.Pow} => (; id = (nothing, nothing), hash2 = 0)
         _ => throw(UnimplementedForVariantError(override_properties, obj))
     end
 end
 
 function ordered_override_properties(obj::Type{<:BSImpl.Variant})
     @match obj begin
-        ::Type{<:BSImpl.Sym} => (0, nothing)
-        ::Type{<:BSImpl.Term} => (0, 0, nothing)
-        ::Type{<:BSImpl.AddOrMul} => (ArgsT(), 0, 0, nothing)
-        ::Type{<:BSImpl.Div} => (0, nothing)
-        ::Type{<:BSImpl.Pow} => (0, nothing)
+        ::Type{<:BSImpl.Sym} => (0, (nothing, nothing))
+        ::Type{<:BSImpl.Term} => (0, 0, (nothing, nothing))
+        ::Type{<:BSImpl.AddOrMul} => (ArgsT(), 0, 0, (nothing, nothing))
+        ::Type{<:BSImpl.Div} => (0, (nothing, nothing))
+        ::Type{<:BSImpl.Pow} => (0, (nothing, nothing))
         _ => throw(UnimplementedForVariantError(override_properties, obj))
     end
 end
@@ -410,8 +410,8 @@ end
 
 function isequal_bsimpl(a::BSImpl.Type, b::BSImpl.Type, full)
     a === b && return true
-    ida = a.id
-    idb = b.id
+    taskida, ida = a.id
+    taskidb, idb = b.id
     ida === idb && ida !== nothing && return true
     typeof(a) === typeof(b) || return false
 
@@ -420,7 +420,7 @@ function isequal_bsimpl(a::BSImpl.Type, b::BSImpl.Type, full)
     Ta === Tb || return false
 
 
-    if full && ida !== idb && ida !== nothing && idb !== nothing
+    if full && ida !== idb && ida !== nothing && idb !== nothing && taskida == taskidb
         return false
     end
 
@@ -596,9 +596,10 @@ const ENABLE_HASHCONSING = Ref(true)
 const WKD = TaskLocalValue{WeakKeyDict{BSImpl.Type, Nothing}}(WeakKeyDict{BSImpl.Type, Nothing})
 const WVD = TaskLocalValue{WeakValueDict{UInt, BSImpl.Type}}(WeakValueDict{UInt, BSImpl.Type})
 const WCS = TaskLocalValue{WeakCacheSet{BSImpl.Type}}(WeakCacheSet{BSImpl.Type})
+const TASK_ID = TaskLocalValue{UInt}(() -> rand(UInt))
 
 function generate_id()
-    return IDType()
+    return (TASK_ID[], IDType())
 end
 
 const TOTAL = TaskLocalValue{Int}(Returns(0))
@@ -651,7 +652,7 @@ function hashcons(s::BSImpl.Type{T})::BSImpl.Type{T} where {T}
         #     cache[h] = s
         #     k = s
         # end
-        if k.id === nothing
+        if k.id === (nothing, nothing)
             k.id = generate_id()
         end
         return k

test/basics.jl

@@ -448,3 +448,13 @@ end
         @test isequal(res, x)
     end
 end
+
+@testset "Equivalent expressions across tasks are equal" begin
+    @syms a
+    task = Threads.@spawn @syms a
+    a2 = only(fetch(task))
+    @test isequal(a, a2)
+    @test SymbolicUtils.@manually_scope SymbolicUtils.COMPARE_FULL => true begin
+        isequal(a, a2)
+    end
+end

test/hash_consing.jl

@@ -129,7 +129,7 @@ end
     @syms a b
     x1 = a + b
     x2 = a + b
-    @test x1.id === nothing === x2.id
+    @test x1.id === (nothing, nothing) === x2.id
     SymbolicUtils.ENABLE_HASHCONSING[] = true
 end