Skip to content

Commit bafe3a8

Browse files
committed
Make Lazy caches support get_tmp
1 parent de6e393 commit bafe3a8

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

src/PreallocationTools.jl

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -216,14 +216,16 @@ function similar_type(x::AbstractArray{T}, s::NTuple{N, Integer}) where {T, N}
216216
typeof(similar(x, ntuple(Returns(1), N)))
217217
end
218218

219-
# override the [] method
220-
function Base.getindex(b::LazyBufferCache, u::T) where {T <: AbstractArray}
219+
function get_tmp(b::LazyBufferCache, u::T) where {T <: AbstractArray}
221220
s = b.sizemap(size(u)) # required buffer size
222221
get!(b.bufs, (T, s)) do
223222
similar(u, s) # buffer to allocate if it was not found in b.bufs
224223
end::similar_type(u, s) # declare type since b.bufs dictionary is untyped
225224
end
226225

226+
# override the [] method
227+
Base.getindex(b::LazyBufferCache, u::T) where {T <: AbstractArray} = get_tmp(b, u)
228+
227229
# GeneralLazyBufferCache
228230

229231
"""
@@ -246,11 +248,12 @@ struct GeneralLazyBufferCache{F <: Function}
246248
GeneralLazyBufferCache(f::F = identity) where {F <: Function} = new{F}(Dict(), f) # start with empty dict
247249
end
248250

249-
function Base.getindex(b::GeneralLazyBufferCache, u::T) where {T}
251+
function get_tmp(b::GeneralLazyBufferCache, u::T) where {T}
250252
get!(b.bufs, T) do
251253
b.f(u)
252254
end
253255
end
256+
Base.getindex(b::GeneralLazyBufferCache, u::T) where {T} = get_tmp(b, u)
254257

255258
export GeneralLazyBufferCache, FixedSizeDiffCache, DiffCache, LazyBufferCache, dualcache
256259
export get_tmp

0 commit comments

Comments
 (0)