Add optional initializer! function to LazyBufferCache

Author: SouthEndMusic

README.md

@@ -218,6 +218,9 @@ A `LazyBufferCache` is a `Dict`-like type for the caches which automatically def
 new cache arrays on demand when they are required. The function `f` maps
 `size_of_cache = f(size(u))`, which by default creates cache arrays of the same size.
 
+By default the created buffers are not initialized, but a function `initializer!`
+can be supplied which is applied to the buffer when it is created, for instance `buf -> fill!(buf, 0.0)`.
+
 Note that `LazyBufferCache` is type-stable and contains no dynamic dispatch. This gives
 it a ~15ns overhead. The upside of `LazyBufferCache` is that the user does not have to
 worry about potential issues with chunk sizes and such: `LazyBufferCache` is much easier!

docs/src/index.md

@@ -208,6 +208,9 @@ A `LazyBufferCache` is a `Dict`-like type for the caches, which automatically de
 new cache arrays on demand when they are required. The function `f` maps
 `size_of_cache = f(size(u))`, which by default creates cache arrays of the same size.
 
+By default the created buffers are not initialized, but a function `initializer!`
+can be supplied which is applied to the buffer when it is created, for instance `buf -> fill!(buf, 0.0)`.
+
 Note that `LazyBufferCache` is type-stable and contains no dynamic dispatch. This gives
 it a ~15ns overhead. The upside of `LazyBufferCache` is that the user does not have to
 worry about potential issues with chunk sizes and such: `LazyBufferCache` is much easier!

src/PreallocationTools.jl

@@ -197,19 +197,27 @@ end
 # LazyBufferCache
 
 """
-    b = LazyBufferCache(f=identity)
+    b = LazyBufferCache(f = identity; initializer! = identity)
 
 A lazily allocated buffer object.  Given an array `u`, `b[u]` returns an array of the
 same type and size `f(size(u))` (defaulting to the same size), which is allocated as
 needed and then cached within `b` for subsequent usage.
 
+By default the created buffers are not initialized, but a function `initializer!`
+can be supplied which is applied to the buffer when it is created, for instance `buf -> fill!(buf, 0.0)`.
+
 Optionally, the size can be explicitly given at calltime using `b[u,s]`, which will
 return a cache of size `s`.
 """
-struct LazyBufferCache{F <: Function}
+struct LazyBufferCache{F <: Function, I <: Function}
     bufs::Dict{Any, Any} # a dictionary mapping (type, size) pairs to buffers
     sizemap::F
-    LazyBufferCache(f::F = identity) where {F <: Function} = new{F}(Dict(), f) # start with empty dict
+    initializer!::I
+    function LazyBufferCache(
+            f::F = identity; initializer!::I = identity) where {
+            F <: Function, I <: Function}
+        new{F, I}(Dict(), f, initializer!)
+    end # start with empty dict
 end
 
 similar_type(x::AbstractArray, s::Integer) = similar_type(x, (s,))
@@ -222,7 +230,9 @@ end
 function get_tmp(
         b::LazyBufferCache, u::T, s = b.sizemap(size(u))) where {T <: AbstractArray}
     get!(b.bufs, (T, s)) do
-        similar(u, s) # buffer to allocate if it was not found in b.bufs
+        buffer = similar(u, s) # buffer to allocate if it was not found in b.bufs
+        b.initializer!(buffer)
+        buffer
     end::similar_type(u, s) # declare type since b.bufs dictionary is untyped
 end
 

test/lbc.jl

@@ -0,0 +1,6 @@
+using PreallocationTools: LazyBufferCache
+using Test
+
+b = LazyBufferCache(Returns(10); initializer! = buf -> fill!(buf, 0))
+
+@test b[Float64[]] == zeros(10)

test/runtests.jl

@@ -16,6 +16,7 @@ if GROUP == "All" || GROUP == "Core"
     @safetestset "DiffCache Resizing" include("core_resizing.jl")
     @safetestset "DiffCache Nested Duals" include("core_nesteddual.jl")
     @safetestset "DiffCache Sparsity Support" include("sparsity_support.jl")
+    @safetestset "LazyBufferCache" include("lbc.jl")
     @safetestset "GeneralLazyBufferCache" include("general_lbc.jl")
 end