Merge pull request #711 from AayushSabharwal/as/fix-sorted-arguments

fix: fix caching breakages

Author: ChrisRackauckas

.github/workflows/Downstream.yml

@@ -23,6 +23,7 @@ jobs:
           - {user: SciML, repo: NeuralPDE.jl, group: NNPDE}
           - {user: SciML, repo: DataDrivenDiffEq.jl, group: Core}
           - {user: JuliaSymbolics, repo: Symbolics.jl, group: All}
+          - {user: JuliaSymbolics, repo: Symbolics.jl, group: Downstream}
           - {user: SciML, repo: ModelOrderReduction.jl, group: All}
 
     steps:

src/SymbolicUtils.jl

@@ -24,6 +24,7 @@ using WeakValueDicts: WeakValueDict
 import ExproniconLite as EL
 import TaskLocalValues: TaskLocalValue
 
+include("cache.jl")
 Base.@deprecate istree iscall
 export istree, operation, arguments, sorted_arguments, iscall
 # Sym, Term,
@@ -76,6 +77,5 @@ include("code.jl")
 # Adjoints
 include("adjoints.jl")
 
-include("cache.jl")
 
 end # module

src/cache.jl

@@ -215,7 +215,7 @@ macro cache(args...)
     fn = EL.JLFunction(fnexpr)
     name = fn.name
     # this will now be an inner workhorse function, which the cached function will call
-    fn.name = gensym(name)
+    fn.name = gensym(Symbol(name))
     # the name of the global constant cache
     cachename = Symbol("cacheof($name)")
     # conditions for performing caching. At the very least, need hashconsing enabled and

src/types.jl

@@ -18,7 +18,6 @@ sdict(kv...) = Dict{Any, Any}(kv...)
 using Base: RefValue
 const EMPTY_ARGS = []
 const EMPTY_HASH = RefValue(UInt(0))
-const NOT_SORTED = RefValue(false)
 const EMPTY_DICT = sdict()
 const EMPTY_DICT_T = typeof(EMPTY_DICT)
 const ENABLE_HASHCONSING = Ref(true)
@@ -42,15 +41,13 @@ const ENABLE_HASHCONSING = Ref(true)
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
         arguments::Vector{Any} = EMPTY_ARGS
-        issorted::RefValue{Bool} = NOT_SORTED
     end
     mutable struct Add{T} <: BasicSymbolic{T}
         coeff::Any             = 0         # exp/den if Pow
         dict::EMPTY_DICT_T     = EMPTY_DICT
         hash::RefValue{UInt}   = EMPTY_HASH
         hash2::RefValue{UInt} = EMPTY_HASH
         arguments::Vector{Any} = EMPTY_ARGS
-        issorted::RefValue{Bool} = NOT_SORTED
     end
     mutable struct Div{T} <: BasicSymbolic{T}
         num::Any               = 1
@@ -150,25 +147,19 @@ end
 
 @inline head(x::BasicSymbolic) = operation(x)
 
-function TermInterface.sorted_arguments(x::BasicSymbolic)
-    args = arguments(x)
+@cache function TermInterface.sorted_arguments(x::BasicSymbolic)::Vector{Any}
+    args = copy(arguments(x))
     @compactified x::BasicSymbolic begin
         Add => @goto ADD
         Mul => @goto MUL
         _   => return args
     end
     @label MUL
-    if !x.issorted[]
-        sort!(args, by=get_degrees)
-        x.issorted[] = true
-    end
+    sort!(args, by=get_degrees)
     return args
 
     @label ADD
-    if !x.issorted[]
-        sort!(args, lt = monomial_lt, by=get_degrees)
-        x.issorted[] = true
-    end
+    sort!(args, lt = monomial_lt, by=get_degrees)
     return args
 end
 
@@ -603,7 +594,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=[], issorted=RefValue(false), kw...)
+    s = Add{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=[], kw...)
     BasicSymbolic(s)
 end
 
@@ -621,7 +612,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=[], issorted=RefValue(false), kw...)
+        s = Mul{T}(; coeff, dict, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), metadata, arguments=[], kw...)
         BasicSymbolic(s)
     end
 end
@@ -642,7 +633,7 @@ ratio(x::Rat,y::Rat) = x//y
 function maybe_intcoeff(x)
     if ismul(x)
         if x.coeff isa Rational && isone(x.coeff.den)
-            Mul{symtype(x)}(; coeff=x.coeff.num, dict=x.dict, x.metadata, arguments=[], issorted=RefValue(false))
+            Mul{symtype(x)}(; coeff=x.coeff.num, dict=x.dict, x.metadata, arguments=[])
         else
             x
         end
@@ -1011,7 +1002,7 @@ end
 function remove_minus(t)
     !iscall(t) && return -t
     @assert operation(t) == (*)
-    args = arguments(t)
+    args = sorted_arguments(t)
     @assert args[1] < 0
     Any[-args[1], args[2:end]...]
 end

test/cache_macro.jl

@@ -2,6 +2,7 @@ using SymbolicUtils
 using SymbolicUtils: BasicSymbolic, @cache, associated_cache, set_limit!, get_limit,
                      clear_cache!, SymbolicKey, metadata, maketerm
 using OhMyThreads: tmap
+using Random
 
 @cache function f1(x::BasicSymbolic)::BasicSymbolic
     return 2x + 1
@@ -29,6 +30,8 @@ end
 
     set_limit!(f1, 10)
     @test get_limit(f1) == 10
+    SymbolicUtils.set_retain_fraction!(f1, 0.1)
+    @test SymbolicUtils.get_retain_fraction(f1) == 0.1
     for i in 1:8
         xx = setmetadata(xx, Int, i)
         f1(xx)
@@ -55,8 +58,6 @@ end
     @test length(cache) == 0
     stats = SymbolicUtils.get_stats(f1)
     @test stats.hits == stats.misses == stats.clears == 0
-    SymbolicUtils.set_retain_fraction!(f1, 0.1)
-    @test SymbolicUtils.get_retain_fraction(f1) == 0.1
     @test SymbolicUtils.is_caching_enabled(f1)
     SymbolicUtils.toggle_caching!(f1, false)
     @test !SymbolicUtils.is_caching_enabled(f1)
@@ -129,7 +130,7 @@ end
     @test stats.hits == stats.misses == stats.clears == 0
 end
 
-@cache function f4(x::Union{BasicSymbolic, Int})::Union{BasicSymbolic, Int}
+@cache function f4(x::Union{BasicSymbolic, Number})::Union{BasicSymbolic, Number}
     x isa Number && return x
     if iscall(x)
         return maketerm(typeof(x), operation(x), map(f4, arguments(x)), metadata(x))
@@ -140,17 +141,17 @@ end
 @testset "Threading" begin
     @syms x y z
     @test isequal(f4(2x + 1), 2(2x + 1) + 1)
-
+    rng = Xoshiro(3)
     function build_rand_expr(vars, depth, maxdepth)
         if depth < maxdepth
             v = build_rand_expr(vars, depth + 1, maxdepth)
         else
-            v = rand(vars)
+            v = rand(rng, vars)
         end
         if isodd(depth)
-            return v + rand([1:3; vars])
+            return v + rand(rng, [1:3; vars])
         else
-            return v * rand([1:3; vars])
+            return v * rand(rng, [1:3; vars])
         end
     end