diff --git a/Project.toml b/Project.toml index c72c3c1..c7c85f4 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/docs/src/index.md b/docs/src/index.md index 22f662f..60fa425 100644 --- a/docs/src/index.md +++ b/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) diff --git a/ext/MemoizationKitSQLiteExt.jl b/ext/MemoizationKitSQLiteExt.jl index ca928c4..845b0fc 100644 --- a/ext/MemoizationKitSQLiteExt.jl +++ b/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) diff --git a/src/call.jl b/src/call.jl index d07349f..91f0e33 100644 --- a/src/call.jl +++ b/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,13 +36,17 @@ 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 @@ -49,15 +54,18 @@ _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) diff --git a/src/disk.jl b/src/disk.jl index 01a85b8..52a6473 100644 --- a/src/disk.jl +++ b/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 diff --git a/src/macro.jl b/src/macro.jl index ea81257..30a95f6 100644 --- a/src/macro.jl +++ b/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)))) diff --git a/test/test_macro.jl b/test/test_macro.jl index 000ad56..cffac62 100644 --- a/test/test_macro.jl +++ b/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