use macro and TLV for scoping

Author: AayushSabharwal

Project.toml

@@ -17,6 +17,7 @@ EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 ExproniconLite = "55351af7-c7e9-48d6-89ff-24e801d99491"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
@@ -54,6 +55,7 @@ DynamicPolynomials = "0.5, 0.6"
 EnumX = "1.0.5"
 ExproniconLite = "0.10.14"
 LabelledArrays = "1.5"
+MacroTools = "0.5.16"
 Moshi = "0.3.6"
 MultivariatePolynomials = "0.5"
 NaNMath = "0.3, 1.1.2"
@@ -89,7 +91,3 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
 test = ["BenchmarkTools", "Documenter", "LabelledArrays", "Pkg", "PkgBenchmark", "Random", "ReferenceTests", "ReverseDiff", "SafeTestsets", "Test", "Zygote", "OhMyThreads", "RuntimeGeneratedFunctions"]
-
-[sources]
-WeakCacheSets = {url="https://github.com/JuliaCollections/WeakCacheSets.jl"}
-Moshi = {url="https://github.com/AayushSabharwal/Moshi.jl", rev="as/mutable-adt"}

src/SymbolicUtils.jl

@@ -30,6 +30,7 @@ import TaskLocalValues: TaskLocalValue
 using WeakValueDicts: WeakValueDict
 using WeakCacheSets: WeakCacheSet, getkey!
 using Base: RefValue
+import MacroTools
 
 # include("WeakCacheSets.jl")
 

src/types.jl

@@ -3,6 +3,7 @@
 #### Symbolic
 #--------------------
 abstract type Symbolic{T} end
+using Setfield: MacroTools
 
 #################### SafeReal #########################
 export SafeReal, LiteralReal
@@ -404,6 +405,53 @@ end
 using Base.ScopedValues
 
 const SV_COMPARE = ScopedValue{Int}()
+const COMPARE_TYPE = TaskLocalValue{Int}(Returns(0))
+
+macro manually_scope(val, expr, is_forced = false)
+    @assert Meta.isexpr(val, :call)
+    @assert val.args[1] == :(=>)
+
+    var_name = val.args[2]
+    new_val = val.args[3]
+    old_name = gensym(:old_val)
+    cur_name = gensym(:cur_val)
+    retval_name = gensym(:retval)
+    close_expr = :($var_name[] = $old_name)
+    interpolated_expr = MacroTools.postwalk(expr) do ex
+        if Meta.isexpr(ex, :return)
+            return Expr(:block, close_expr, ex)
+        elseif Meta.isexpr(ex, :$) && length(ex.args) == 1 && ex.args[1] == :$
+            return cur_name
+        else
+            return ex
+        end
+    end
+    basic_result = quote
+        $cur_name = $var_name[] = $new_val
+        $retval_name = begin
+            $interpolated_expr
+        end
+        $close_expr
+        $retval_name
+    end
+    is_forced && return quote
+        $old_name = $var_name[]
+        $basic_result
+    end |> esc
+
+    return quote
+        $old_name = $var_name[]
+        if $iszero($old_name)
+            $basic_result
+        else
+            $cur_name = $old_name
+            $retval_name = begin
+                $interpolated_expr
+            end
+        end
+        $retval_name
+    end |> esc
+end
 
 function isequal_symdict(a::Dict, b::Dict, val)
     if val == 2
@@ -413,11 +461,11 @@ function isequal_symdict(a::Dict, b::Dict, val)
     for (k, v) in a
         k2 = nothing
         v2 = nothing
-        @with SV_COMPARE => 2 begin
+        @manually_scope COMPARE_TYPE => 2 begin
             k2 = getkey(b, k, nothing)
             k2 === nothing && return false
             v2 = b[k2]
-        end
+        end true
         v == v2 && isequal(k, k2) || return false
     end
     return true
@@ -462,13 +510,8 @@ function Base.isequal(a::BSImpl.Type, b::BSImpl.Type)
     Tb = MData.variant_type(b)
     Ta === Tb || return false
 
-    val = ScopedValues.get(SV_COMPARE)
-    if val === nothing
-        @with SV_COMPARE => 1 begin
-            isequal_bsimpl(a, b, 1)
-        end
-    else
-        isequal_bsimpl(a, b, something(val))
+    @manually_scope COMPARE_TYPE => 1 begin
+        isequal_bsimpl(a, b, $$)
     end
 end
 
@@ -477,12 +520,7 @@ function Base.isequal(a::BasicSymbolic, b::BasicSymbolic)
     typeof(a) === typeof(b) || return false
 
 
-    val = ScopedValues.get(SV_COMPARE)
-    if val === nothing
-        @with SV_COMPARE => 2 begin
-            isequal(_unwrap_internal(a), _unwrap_internal(b))
-        end
-    else
+    @manually_scope COMPARE_TYPE => 2 begin
         isequal(_unwrap_internal(a), _unwrap_internal(b))
     end
 end
@@ -492,12 +530,7 @@ end
 for T1 in [BasicSymbolic, BSImpl.Type], T2 in [BasicSymbolic, BSImpl.Type]
     T1 == T2 && continue
     @eval function Base.isequal(a::$T1, b::$T2)
-        val = ScopedValues.get(SV_COMPARE)
-        if val === nothing
-            @with SV_COMPARE => 2 begin
-                isequal(_unwrap_internal(a), _unwrap_internal(b))
-            end
-        else
+        @manually_scope COMPARE_TYPE => 2 begin
             isequal(_unwrap_internal(a), _unwrap_internal(b))
         end
     end
@@ -639,30 +672,14 @@ function Base.hash(s::BSImpl.Type, h::UInt)
     if !iszero(h)
         return hash(hash(s, zero(h)), h)::UInt
     end
-    val = ScopedValues.get(SV_COMPARE)
-    if val === nothing
-        @with SV_COMPARE => 1 begin
-            hash_bsimpl(s, h, 1)
-        end
-    else
-        hash_bsimpl(s, h, something(val))
+    @manually_scope COMPARE_TYPE => 1 begin
+        hash_bsimpl(s, h, $$)
     end
 end
 
 Base.@nospecializeinfer function Base.hash(x::BasicSymbolic, h::UInt)
     @nospecialize x
-    val = ScopedValues.get(SV_COMPARE)
-    if val === nothing
-        @with SV_COMPARE => 2 begin
-            if x isa BasicSymbolic{Real}
-                result = Base.hash(_unwrap_internal(x), h)
-            elseif x isa BasicSymbolic{Number}
-                result = Base.hash(_unwrap_internal(x), h)
-            else
-                result = Base.hash(_unwrap_internal(x), h)
-            end
-        end
-    else
+    @manually_scope COMPARE_TYPE => 2 begin
         if x isa BasicSymbolic{Real}
             result = Base.hash(_unwrap_internal(x), h)
         elseif x isa BasicSymbolic{Number}