Merge pull request #709 from AayushSabharwal/as/cache-fix

fix: handle explicit `::Any` annotation in `@cache` macro

Author: AayushSabharwal

src/cache.jl

@@ -229,6 +229,14 @@ macro cache(args...)
     keytypes = []
     # The arguments of the workhorse function
     argexprs = []
+    # The name of the variable storing the result of looking up the cache
+    cache_value_name = :val
+    # The condition for a cache hit
+    cache_hit_condition = :(!($cache_value_name isa $CacheSentinel))
+    # Type of additional data stored with cached result. Used to compare
+    # equality of `BasicSymbolic` arguments, since `objectid` is a hash.
+    cache_additional_types = []
+    cache_additional_values = []
 
     for arg in fn.args
         # handle arguments with defaults
@@ -240,18 +248,39 @@ macro cache(args...)
             push!(keyexprs, :($arg isa BasicSymbolic ? $SymbolicKey(objectid($arg)) : $arg))
             push!(argexprs, arg)
             push!(keytypes, Any)
+            push!(cache_additional_types, Any)
+            push!(cache_additional_values, arg)
+            cache_hit_condition = :($cache_hit_condition && (!($arg isa BasicSymbolic) || $arg === $cache_value_name[$(length(cache_additional_values))]))
             continue
         end
         argname, Texpr = arg.args
         push!(argexprs, argname)
 
+        if Texpr == :Any
+            # if the type is `Any`, branch on it being a `BasicSymbolic`
+            push!(keyexprs, :($argname isa BasicSymbolic ? $SymbolicKey(objectid($argname)) : $argname))
+            push!(keytypes, Any)
+            push!(cache_additional_types, Any)
+            push!(cache_additional_values, argname)
+            cache_hit_condition = :($cache_hit_condition && (!($argname isa BasicSymbolic) || $argname === $cache_value_name[$(length(cache_additional_values))]))
+            continue
+        end
+
         # handle Union types that may contain a `BasicSymbolic`
         if Meta.isexpr(Texpr, :curly) && Texpr.args[1] == :Union
             Texprs = Texpr.args[2:end]
             Ts = map(Base.Fix1(Base.eval, __module__), Texprs)
             keyTs = map(x -> x <: BasicSymbolic ? SymbolicKey : x, Ts)
+            maybe_basicsymbolic = any(x -> x <: BasicSymbolic, Ts)
             push!(keytypes, Union{keyTs...})
-            push!(keyexprs, :($argname isa BasicSymbolic ? $SymbolicKey(objectid($argname)) : $argname))
+            if maybe_basicsymbolic
+                push!(keyexprs, :($argname isa BasicSymbolic ? $SymbolicKey(objectid($argname)) : $argname))
+                push!(cache_additional_types, Texpr)
+                push!(cache_additional_values, argname)
+                cache_hit_condition = :($cache_hit_condition && (!($argname isa BasicSymbolic) || $argname === $cache_value_name[$(length(cache_additional_values))]))
+            else
+                push!(keyexprs, argname)
+            end
             continue
         end
 
@@ -260,6 +289,9 @@ macro cache(args...)
         if T <: BasicSymbolic
             push!(keytypes, SymbolicKey) 
             push!(keyexprs, :($SymbolicKey(objectid($argname))))
+            push!(cache_additional_types, T)
+            push!(cache_additional_values, argname)
+            cache_hit_condition = :($cache_hit_condition && $argname === $cache_value_name[$(length(cache_additional_values))])
         else
             push!(keytypes, T)
             push!(keyexprs, argname)
@@ -280,8 +312,11 @@ macro cache(args...)
     # construct an expression for the type of the cache keys
     keyT = Expr(:curly, Tuple)
     append!(keyT.args, keytypes)
+    valT = Expr(:curly, Tuple)
+    append!(valT.args, cache_additional_types)
+    push!(valT.args, rettype)
     # the type of the cache
-    cacheT = :(Dict{$keyT, $rettype})
+    cacheT = :(Dict{$keyT, $valT})
     # type of the `TaskLocalValue`
     tlvT = :($(TaskLocalValue){Tuple{$cacheT, $CacheStats}})
     # the name of the cache struct
@@ -324,11 +359,11 @@ macro cache(args...)
             # look it up
             # we use a custom sentinel value since `nothing` is a valid return value
             # which we might want to cache
-            val = $(get)(cachedict, key, $(CacheSentinel)())
-            if !(val isa $CacheSentinel)
+            $cache_value_name = $(get)(cachedict, key, $(CacheSentinel)())
+            if $cache_hit_condition
                 # cache hit
                 cachestats.hits += 1
-                return val
+                return $cache_value_name[end]
             end
             # cache miss
             cachestats.misses += 1
@@ -339,7 +374,7 @@ macro cache(args...)
                 $(filter!)($cachename, cachedict)
             end
             # add to cache
-            cachedict[key] = val
+            cachedict[key] = ($(cache_additional_values...), val)
             return val
         end
 

src/types.jl

@@ -69,6 +69,14 @@ function SymbolicIndexingInterface.symbolic_type(::Type{<:BasicSymbolic})
     ScalarSymbolic()
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Return the inner `Symbolic` wrapped in a non-symbolic subtype. Defaults to
+returning the input as-is.
+"""
+unwrap(x) = x
+
 function exprtype(x::BasicSymbolic)
     @compactified x::BasicSymbolic begin
         Term => TERM
@@ -538,10 +546,17 @@ function Sym{T}(name::Symbol; kw...) where {T}
     BasicSymbolic(s)
 end
 
+function unwrap_arr!(arr)
+    for i in eachindex(arr)
+        arr[i] = unwrap(arr[i])
+    end
+end
+
 function Term{T}(f, args; kw...) where T
     if eltype(args) !== Any
         args = convert(Vector{Any}, args)
     end
+    unwrap_arr!(args)
 
     s = Term{T}(;f=f, arguments=args, hash=Ref(UInt(0)), hash2=Ref(UInt(0)), kw...)
     BasicSymbolic(s)
@@ -551,7 +566,16 @@ function Term(f, args; metadata=NO_METADATA)
     Term{_promote_symtype(f, args)}(f, args, metadata=metadata)
 end
 
+function unwrap_dict(dict)
+    if any(k -> unwrap(k) !== k, keys(dict))
+        return typeof(dict)(unwrap(k) => v for (k, v) in dict)
+    end
+    return dict
+end
+
 function Add(::Type{T}, coeff, dict; metadata=NO_METADATA, kw...) where T
+    coeff = unwrap(coeff)
+    dict = unwrap_dict(dict)
     if isempty(dict)
         return coeff
     elseif _iszero(coeff) && length(dict) == 1
@@ -569,6 +593,8 @@ function Add(::Type{T}, coeff, dict; metadata=NO_METADATA, kw...) where T
 end
 
 function Mul(T, a, b; metadata=NO_METADATA, kw...)
+    a = unwrap(a)
+    b = unwrap_dict(b)
     isempty(b) && return a
     if _isone(a) && length(b) == 1
         pair = first(b)
@@ -613,6 +639,8 @@ function maybe_intcoeff(x)
 end
 
 function Div{T}(n, d, simplified=false; metadata=nothing, kwargs...) where {T}
+    n = unwrap(n)
+    d = unwrap(d)
     if T<:Number && !(T<:SafeReal)
         n, d = quick_cancel(n, d)
     end
@@ -662,6 +690,8 @@ end
 @inline denominators(x) = isdiv(x) ? numerators(x.den) : Any[1]
 
 function Pow{T}(a, b; metadata=NO_METADATA, kwargs...) where {T}
+    a = unwrap(a)
+    b = unwrap(b)
     _iszero(b) && return 1
     _isone(b) && return a
     s = Pow{T}(; base=a, exp=b, arguments=[], metadata)
@@ -760,6 +790,7 @@ function makepow(a, b)
 end
 
 function term(f, args...; type = nothing)
+    args = map(unwrap, args)
     if type === nothing
         T = _promote_symtype(f, args)
     else
@@ -894,6 +925,7 @@ Base.ImmutableDict(d::ImmutableDict{K,V}, x, y)  where {K, V} = ImmutableDict{K,
 
 assocmeta(d::Dict, ctx, val) = (d=copy(d); d[ctx] = val; d)
 function assocmeta(d::Base.ImmutableDict, ctx, val)::ImmutableDict{DataType,Any}
+    val = unwrap(val)
     # optimizations
     # If using upto 3 contexts, things stay compact
     if isdefined(d, :parent)
@@ -915,7 +947,7 @@ function setmetadata(s::Symbolic, ctx::DataType, val)
         @set s.metadata = assocmeta(s.metadata, ctx, val)
     else
         # fresh Dict
-        @set s.metadata = Base.ImmutableDict{DataType, Any}(ctx, val)
+        @set s.metadata = Base.ImmutableDict{DataType, Any}(ctx, unwrap(val))
     end
 end
 
@@ -1333,6 +1365,10 @@ function +(a::SN, b::SN)
 end
 
 function +(a::Number, b::SN)
+    tmp = unwrap(a)
+    if tmp !== a
+        return tmp + b
+    end
     !issafecanon(+, b) && return term(+, a, b) # Don't flatten if args have metadata
     iszero(a) && return b
     if isadd(b)
@@ -1399,6 +1435,10 @@ function *(a::SN, b::SN)
 end
 
 function *(a::Number, b::SN)
+    tmp = unwrap(a)
+    if tmp !== a
+        return tmp * b
+    end
     !issafecanon(*, b) && return term(*, a, b)
     if iszero(a)
         a
@@ -1439,6 +1479,7 @@ end
 ###
 
 function ^(a::SN, b)
+    b = unwrap(b)
     !issafecanon(^, a,b) && return Pow(a, b)
     if b isa Number && iszero(b)
         # fast path

test/cache_macro.jl

@@ -13,9 +13,9 @@ end
     @test isequal(val, 2x + 1)
     cachestruct = associated_cache(f1)
     cache, stats = cachestruct.tlv[]
-    @test cache isa Dict{Tuple{SymbolicKey}, BasicSymbolic}
+    @test cache isa Dict{Tuple{SymbolicKey}, Tuple{BasicSymbolic, BasicSymbolic}}
     @test length(cache) == 1
-    @test cache[(SymbolicKey(objectid(x)),)] === val
+    @test cache[(SymbolicKey(objectid(x)),)][end] === val
     @test stats.hits == 0
     @test stats.misses == 1
     f1(x)
@@ -75,9 +75,9 @@ end
     @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 cache isa Dict{Tuple{Union{SymbolicKey, UInt}}, NTuple{2, Union{BasicSymbolic, UInt}}}
     @test length(cache) == 1
-    @test cache[(SymbolicKey(objectid(x)),)] === val
+    @test cache[(SymbolicKey(objectid(x)),)][end] === val
     @test stats.hits == 0
     @test stats.misses == 1
     f2(x)
@@ -88,7 +88,7 @@ end
     val = f2(y)
     @test val == 2y + 1
     @test length(cache) == 2
-    @test cache[(y,)] == val
+    @test cache[(y,)][end] == val
     @test stats.misses == 2
 
     clear_cache!(f2)
@@ -100,27 +100,31 @@ end
     return 2x + 1
 end
 
-@testset "::Any" begin
+@cache function f3_2(x::Any)::Union{BasicSymbolic, Int}
+    return 2x + 1
+end
+
+@testset "$name" for (name, fn) in [("implicit ::Any", f3), ("explicit ::Any", f3_2)]
     @syms x
-    val = f3(x)
+    val = fn(x)
     @test isequal(val, 2x + 1)
-    cachestruct = associated_cache(f3)
+    cachestruct = associated_cache(fn)
     cache, stats = cachestruct.tlv[]
-    @test cache isa Dict{Tuple{Any}, Union{BasicSymbolic, Int}}
+    @test cache isa Dict{Tuple{Any}, Tuple{Any, Union{BasicSymbolic, Int}}}
     @test length(cache) == 1
-    @test cache[(SymbolicKey(objectid(x)),)] === val
+    @test cache[(SymbolicKey(objectid(x)),)][end] === val
     @test stats.hits == 0
     @test stats.misses == 1
-    f3(x)
+    fn(x)
     @test stats.hits == 1
     @test stats.misses == 1
 
-    val = f3(3)
+    val = fn(3)
     @test val == 7
     @test length(cache) == 2
     @test stats.misses == 2
 
-    clear_cache!(f3)
+    clear_cache!(fn)
     @test length(cache) == 0
     @test stats.hits == stats.misses == stats.clears == 0
 end
@@ -155,3 +159,39 @@ end
     truevals = map(f4, exprs)
     @test isequal(result, truevals)
 end
+
+@cache function f5(x::BasicSymbolic, y::Union{BasicSymbolic, Int}, z)::BasicSymbolic
+    return x + y + z
+end
+
+# temporary defintion to induce objectid collisions
+Base.objectid(x::BasicSymbolic) = 0x42
+
+@testset "`objectid` collision handling" begin
+    @syms x y z
+    @test objectid(x) == objectid(y) == objectid(z) == 0x42
+    cachestruct = associated_cache(f5)
+    cache, stats = cachestruct.tlv[]
+    val = f5(x, 1, 2)
+    @test isequal(val, x + 3)
+    @test length(cache) == 1
+    @test stats.misses == 1
+    val2 = f5(y, 1, 2)
+    @test isequal(val2, y + 3)
+    @test length(cache) == 1
+    @test stats.misses == 2
+
+    clear_cache!(f5)
+    val = f5(x, y, z)
+    @test isequal(val, x + y + z)
+    @test length(cache) == 1
+    @test stats.misses == 1
+    val2 = f5(y, 2z, x)
+    @test isequal(val2, x + y + 2z)
+    @test length(cache) == 1
+    @test stats.misses == 2
+end
+
+Base.delete_method(only(methods(objectid, @__MODULE__)))
+@syms x
+@test objectid(x) != 0x42