test: test @cache macro

Author: AayushSabharwal

Project.toml

@@ -51,6 +51,7 @@ ExproniconLite = "0.10.14"
 LabelledArrays = "1.5"
 MultivariatePolynomials = "0.5"
 NaNMath = "0.3, 1.1.2"
+OhMyThreads = "0.7"
 ReverseDiff = "1"
 Setfield = "0.7, 0.8, 1"
 SpecialFunctions = "0.10, 1.0, 2"
@@ -67,6 +68,7 @@ julia = "1.10"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -77,4 +79,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["BenchmarkTools", "Documenter", "LabelledArrays", "Pkg", "PkgBenchmark", "Random", "ReferenceTests", "ReverseDiff", "SafeTestsets", "Test", "Zygote"]
+test = ["BenchmarkTools", "Documenter", "LabelledArrays", "Pkg", "PkgBenchmark", "Random", "ReferenceTests", "ReverseDiff", "SafeTestsets", "Test", "Zygote", "OhMyThreads"]

test/cache_macro.jl

@@ -0,0 +1,157 @@
+using SymbolicUtils
+using SymbolicUtils: BasicSymbolic, @cache, associated_cache, set_limit!, get_limit,
+                     clear_cache!, SymbolicKey, metadata, maketerm
+using OhMyThreads: tmap
+
+@cache function f1(x::BasicSymbolic)::BasicSymbolic
+    return 2x + 1
+end
+
+@testset "::BasicSymbolic" begin
+    @syms x
+    val = f1(x)
+    @test isequal(val, 2x + 1)
+    cachestruct = associated_cache(f1)
+    cache, stats = cachestruct.tlv[]
+    @test cache isa Dict{Tuple{SymbolicKey}, BasicSymbolic}
+    @test length(cache) == 1
+    @test cache[(SymbolicKey(objectid(x)),)] === val
+    @test stats.hits == 0
+    @test stats.misses == 1
+    f1(x)
+    @test stats.hits == 1
+    @test stats.misses == 1
+
+    xx = setmetadata(x, Int, 0)
+    val = f1(xx)
+    @test length(cache) == 2
+    @test stats.misses == 2
+
+    set_limit!(f1, 10)
+    @test get_limit(f1) == 10
+    for i in 1:8
+        xx = setmetadata(xx, Int, i)
+        f1(xx)
+        @test length(cache) == i + 2
+    end
+    xx = setmetadata(xx, Int, 9)
+    f1(xx)
+    @test length(cache) < 10
+    @test stats.clears == 1
+
+    hits = stats.hits
+    misses = stats.misses
+    len = length(cache)
+
+    @syms x::Float64 # different symtype
+    val = f1(x)
+    @test length(cache) == len + 1
+    @test stats.hits == hits
+    @test stats.misses == misses + 1
+    @test f1(x) === val
+    @test stats.hits == hits + 1
+    
+    clear_cache!(f1)
+    @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)
+    f1(x)
+    @test isempty(cache)
+    @test stats.hits == stats.misses == stats.clears == 0
+end
+
+@cache function f2(x::Union{BasicSymbolic, UInt})::Union{BasicSymbolic, UInt}
+    return 2x + 1
+end
+
+@testset "::Union (with `UInt`)" begin
+    @syms x
+    val = f2(x)
+    @test isequal(val, 2x + 1)
+    cachestruct = associated_cache(f2)
+    cache, stats = cachestruct.tlv[]
+    @test cache isa Dict{Tuple{Union{SymbolicKey, UInt}}, Union{BasicSymbolic, UInt}}
+    @test length(cache) == 1
+    @test cache[(SymbolicKey(objectid(x)),)] === val
+    @test stats.hits == 0
+    @test stats.misses == 1
+    f2(x)
+    @test stats.hits == 1
+    @test stats.misses == 1
+
+    y = objectid(x)
+    val = f2(y)
+    @test val == 2y + 1
+    @test length(cache) == 2
+    @test cache[(y,)] == val
+    @test stats.misses == 2
+
+    clear_cache!(f2)
+    @test length(cache) == 0
+    @test stats.hits == stats.misses == stats.clears == 0
+end
+
+@cache function f3(x)::Union{BasicSymbolic, Int}
+    return 2x + 1
+end
+
+@testset "::Any" begin
+    @syms x
+    val = f3(x)
+    @test isequal(val, 2x + 1)
+    cachestruct = associated_cache(f3)
+    cache, stats = cachestruct.tlv[]
+    @test cache isa Dict{Tuple{Any}, Union{BasicSymbolic, Int}}
+    @test length(cache) == 1
+    @test cache[(SymbolicKey(objectid(x)),)] === val
+    @test stats.hits == 0
+    @test stats.misses == 1
+    f3(x)
+    @test stats.hits == 1
+    @test stats.misses == 1
+
+    val = f3(3)
+    @test val == 7
+    @test length(cache) == 2
+    @test stats.misses == 2
+
+    clear_cache!(f3)
+    @test length(cache) == 0
+    @test stats.hits == stats.misses == stats.clears == 0
+end
+
+@cache function f4(x::Union{BasicSymbolic, Int})::Union{BasicSymbolic, Int}
+    x isa Number && return x
+    if iscall(x)
+        return maketerm(typeof(x), operation(x), map(f4, arguments(x)), metadata(x))
+    end
+    return f3(x)
+end
+
+@testset "Threading" begin
+    @syms x y z
+    @test isequal(f4(2x + 1), 2(2x + 1) + 1)
+
+    function build_rand_expr(vars, depth, maxdepth)
+        if depth < maxdepth
+            v = build_rand_expr(vars, depth + 1, maxdepth)
+        else
+            v = rand(vars)
+        end
+        if isodd(depth)
+            return v + rand([1:3; vars])
+        else
+            return v * rand([1:3; vars])
+        end
+    end
+
+    exprs = [build_rand_expr([x, y, z], 0, 100) for _ in 1:1000]
+    result = tmap(f4, exprs)
+    truevals = map(f4, exprs)
+    @test isequal(result, truevals)
+end

test/runtests.jl

@@ -17,5 +17,6 @@ using Pkg, Test, SafeTestsets
         @safetestset "Fuzz" begin include("fuzz.jl") end
         @safetestset "Adjoints" begin include("adjoints.jl") end
         @safetestset "Hash Consing" begin include("hash_consing.jl") end
+        @safetestset "Cache macro" begin include("cache_macro.jl") end
     end
 end