Apply return annotations in the wrapper and keep Type arguments in inference

Project.toml

@@ -1,6 +1,6 @@
 name = "MemoizationKit"
 uuid = "1b238080-9255-4fe9-b224-89eb24efe93b"
-version = "0.1.1"
+version = "0.1.2"
 authors = ["Lukas Devos <ldevos98@gmail.com> and contributors"]
 
 [workspace]

docs/src/index.md

@@ -60,11 +60,14 @@ Anonymous functions are not supported.
 Methods without `@cached` keep their usual behavior.
 
 A return annotation is optional and may depend on `where` parameters.
+As usual in Julia, the result is converted to the annotated type, before it is cached.
 Without one, MemoizationKit uses the inferred return type, falling back to `Any` if it is not concrete.
+Type arguments keep their own type in this inference, so a body that depends on them stays inferred.
 A type-stable body keeps calls inferred even though the shared cache holds results of different types.
 
 Use [`uncached`](@ref) to call the body without reading or filling RAM or disk caches.
 Supply all arguments, including defaults: defaults belong to the wrapper, not the uncached body.
+Likewise, the return type annotation belongs to the wrapper, so `uncached` returns the value of the body as is.
 
 ```@example powers
 uncached(powers, 2, 3; offset = 1)

ext/MemoizationKitSQLiteExt.jl

@@ -140,14 +140,14 @@ end
 
 function MemoizationKit.disk_lookup(f::F, ::DiskCache{S}, ::Type{V}, key, args, kw, o) where {F, S, V}
     s = stores(f)
-    s.artifact === nothing && s.node === nothing && return _compute(f, o, args, kw)
+    s.artifact === nothing && s.node === nothing && return _compute(f, o, V, args, kw)
     return instrument(f, Val(:disk), () -> lookup(f, s, S, V, key, args, kw, o), o)
 end
 
 function lookup(f::F, s::Stores, ::Type{S}, ::Type{V}, key, args, kw, o) where {F, S, V}
     # callable objects with fields share the store of their type, so they are part of the key
     kb = attempt(() -> tobytes(S, Base.issingletontype(F) ? key : (f, key)), "serialize a key of `$(_fname(f))`")
-    kb === nothing && (Threads.atomic_add!(s.misses, 1); return _compute(f, o, args, kw))
+    kb === nothing && (Threads.atomic_add!(s.misses, 1); return _compute(f, o, V, args, kw))
     h = sha256(kb)
     for store in (s.artifact, s.node)
         store === nothing && continue
@@ -158,7 +158,7 @@ function lookup(f::F, s::Stores, ::Type{S}, ::Type{V}, key, args, kw, o) where {
         v !== nothing && something(v) isa V && (Threads.atomic_add!(s.hits, 1); return something(v)::V)
     end
     Threads.atomic_add!(s.misses, 1)
-    v = _compute(f, o, args, kw)
+    v = _compute(f, o, V, args, kw)
     node = s.node
     node === nothing || attempt("write `$(_fname(f))` to the disk cache at $(node.db.file)") do
         vb = tobytes(S, v)

src/call.jl

@@ -11,6 +11,7 @@ function implementation end
 
 Call the [`@cached`](@ref) function `f` without consulting or filling any cache.
 Supply all positional and keyword arguments, including those with defaults in `f`.
+The result is not converted to the return type annotation of `f`, if there is one.
 """
 uncached(f, args...; kwargs...) = implementation(f, args...; kwargs...)
 
@@ -35,29 +36,36 @@ instrument_label(f, ::Val{phase}) where {phase} = string(phase, ' ', _fname(f))
 
 _callimpl(f, args, kw) = implementation(f, args...; kw...)
 
+# The argument types `A` of a call to an unannotated `@cached` method, as seen by dispatch:
+# `typeof(args)` would turn a `Type{T}` argument into `DataType`, losing `T`.
+struct InferValue{A <: Tuple} end
+
 # Entry point of the methods generated by `@cached`. `vt` is the return type annotation, or
-# `nothing` when there is none, in which case the value type is inferred.
+# `InferValue` when there is none, in which case the value type is inferred.
 # The disk strategy `DiskCacheStyle(f, args...)` is resolved here too; without disk caching, the
 # call is the same as before (`_call`).
 @inline call(f::F, style::CacheStyle, ::Type{V}, args::Tuple, kw::NamedTuple) where {F, V} =
     _call(f, style, DiskCacheStyle(f, args...), V, _callkey(f, args, kw), args, kw)
-@inline call(f::F, style::CacheStyle, ::Nothing, args::A, kw::KW) where {F, A <: Tuple, KW <: NamedTuple} =
+@inline call(f::F, style::CacheStyle, ::InferValue{A}, args::Tuple, kw::KW) where {F, A, KW <: NamedTuple} =
     _call(f, style, DiskCacheStyle(f, args...), _inferred_valtype(Core.Compiler.return_type(_callimpl, Tuple{F, A, KW})), _callkey(f, args, kw), args, kw)
 
 _inferred_valtype(::Type{V}) where {V} = isconcretetype(V) ? V : Any
 
 # `f::F` forces specialization on the function, which Julia skips for `Function` arguments that
 # are only passed along; without it the closures would box `f` and allocate. The owner `o` is
 # computed outside the closures, which Julia 1.10 otherwise boxes in the first caller compiled.
-@inline _compute(f::F, o, args, kw) where {F} = instrument(f, Val(:compute), () -> implementation(f, args...; kw...), o)
+# The result is converted to the value type `V`, which applies a return type annotation of `f`.
+@inline function _compute(f::F, o, ::Type{V}, args, kw)::V where {F, V}
+    return instrument(f, Val(:compute), () -> implementation(f, args...; kw...), o)
+end
 
-_call(f::F, ::NoCache, ::Type{V}, key, args, kw) where {F, V} = _compute(f, _ownerval(F), args, kw)::V
+_call(f::F, ::NoCache, ::Type{V}, key, args, kw) where {F, V} = _compute(f, _ownerval(F), V, args, kw)
 
 # The cache holds keys and values of any type; the assertion recovers the value type, so that
 # calls stay inferred. Lookups do not box the key (see `Key`).
 @inline function _call(f::F, style::CacheStyle, ::Type{V}, key, args, kw) where {F, V}
     o = _ownerval(F)
-    return instrument(f, Val(:lookup), () -> get!(() -> _compute(f, o, args, kw), _container(f, style, key, V), key), o)::V
+    return instrument(f, Val(:lookup), () -> get!(() -> _compute(f, o, V, args, kw), _container(f, style, key, V), key), o)::V
 end
 
 @inline _call(f::F, style::CacheStyle, ::NoCache, ::Type{V}, key, args, kw) where {F, V} = _call(f, style, V, key, args, kw)

src/disk.jl

@@ -66,7 +66,7 @@ enable_disk_caches!() = (DISK_ENABLED[] = true; nothing)
 # Called on a miss of the RAM cache (or on every call, with `NoCache` in RAM).
 function _diskcall(f::F, disk::DiskCache, ::Type{V}, key, args, kw) where {F, V}
     o = _ownerval(F)
-    (DISK_ENABLED[] && ccall(:jl_generating_output, Cint, ()) == 0) || return _compute(f, o, args, kw)
+    (DISK_ENABLED[] && ccall(:jl_generating_output, Cint, ()) == 0) || return _compute(f, o, V, args, kw)
     return disk_lookup(f, disk, V, key, args, kw, o)
 end
 

src/macro.jl

@@ -43,9 +43,15 @@ macro cached(ex)
     impl[:args] = [selfarg; map(_nodefault, args)]
     isempty(kwargs) || (impl[:kwargs] = map(_nodefault, kwargs))
 
-    vt = get(def, :rtype, nothing)
+    vt = get(def, :rtype, :($InferValue{Tuple{$(map(_argtype, args)...)}}()))
     positional = Expr(:tuple, map(_forward, args)...)
     keywords = Expr(:tuple, Expr(:parameters, map(_forward, kwargs)...))
+
+    # The return type annotation is evaluated only in the wrapper, and applied by `call` to the
+    # computed value: in the body, where both depend on the arguments, JET would report the pairs
+    # of value and annotated type that inference cannot rule out.
+    delete!(impl, :rtype)
+    delete!(def, :rtype)
     style = :($CacheStyle($fself, $(positional.args...)))
     def[:body] = :($call($fself, $style, $vt, $positional, $keywords))
 
@@ -72,3 +78,4 @@ _nameargs(a) = Meta.isexpr(a, :(::)) && length(a.args) == 1 ? Expr(:(::), gensym
 
 _nodefault(a) = Meta.isexpr(a, :kw) ? a.args[1] : a
 _forward(a) = Meta.isexpr(a, :(...)) ? Expr(:(...), _argname(a)) : _argname(a)
+_argtype(a) = Meta.isexpr(a, :(...)) ? :(map($Core.Typeof, $(_argname(a)))...) : :($Core.Typeof($(_argname(a))))

test/test_macro.jl

@@ -26,6 +26,10 @@ end
 
 @cached unstable(x) = counting(x > 0 ? x : "negative")
 
+# value types inferred from type arguments, which `typeof` would widen to `DataType`
+@cached typearg(::Type{T}, n::Int) where {T} = counting(Vector{T}(undef, n))
+@cached typevararg(Ts::Type...) = counting(map(T -> Vector{T}(), Ts))
+
 struct Adder
     a::Int
 end
@@ -155,6 +159,7 @@ end
 
 @testset "return types" begin
     @test converted(1) === 1.0
+    @test uncached(converted, 2) === 2 # the annotation applies to cached calls
     @test dependent(1) == [1] && dependent(1) isa Vector{Int}
     @test dependent(1.5) isa Vector{Float64}
     # one untyped cache per function; the value type is recovered at the call
@@ -166,6 +171,8 @@ end
     @test rettype(basic, Tuple{Int}) == Int
     @test unstable(1) === 1 && unstable(-1) == "negative"
     @test rettype(unstable, Tuple{Int}) == Any
+    @test rettype(typearg, Tuple{Type{Float64}, Int}) == Vector{Float64}
+    @test rettype(typevararg, Tuple{Type{Int}, Type{Float64}}) == Tuple{Vector{Int}, Vector{Float64}}
     @test length(only(caches(unstable))) == 2
 end
 
@@ -174,6 +181,8 @@ end
     @test @inferred(full(1, 2, :a; k = 2)) == 5
     @test @inferred(dependent(2)) == [2]
     @test @inferred(Adder(1)(1)) == 2
+    @test @inferred(typearg(Float64, 2)) isa Vector{Float64}
+    @test @inferred(typevararg(Int, Float64)) == (Int[], Float64[])
     basic(6)
     @test (@allocated basic(6)) == 0
 end