Merge pull request #100 from LilithHafner/lh/lbc-fix

Fix and test LazyBufferCache on types that are not fixed points of similar [bugfix]

Author: ChrisRackauckas

docs/src/index.md

@@ -316,7 +316,7 @@ tries to do this with a bump allocator.
 
   - See the [SciML Style Guide](https://github.com/SciML/SciMLStyle) for common coding practices and other style decisions.
   - There are a few community forums:
-
+    
       + The #diffeq-bridged and #sciml-bridged channels in the
         [Julia Slack](https://julialang.org/slack/)
       + The #diffeq-bridged and #sciml-bridged channels in the

src/PreallocationTools.jl

@@ -9,7 +9,7 @@ struct FixedSizeDiffCache{T <: AbstractArray, S <: AbstractArray}
 end
 
 function FixedSizeDiffCache(u::AbstractArray{T}, siz,
-    ::Type{Val{chunk_size}}) where {T, chunk_size}
+        ::Type{Val{chunk_size}}) where {T, chunk_size}
     x = ArrayInterface.restructure(u,
         zeros(ForwardDiff.Dual{nothing, T, chunk_size},
             siz...))
@@ -25,8 +25,8 @@ and for the `Dual` version of `u`, allowing use of pre-cached vectors with
 forward-mode automatic differentiation.
 """
 function FixedSizeDiffCache(u::AbstractArray,
-    ::Type{Val{N}} = Val{ForwardDiff.pickchunksize(length(u))}) where {
-    N,
+        ::Type{Val{N}} = Val{ForwardDiff.pickchunksize(length(u))}) where {
+        N,
 }
     FixedSizeDiffCache(u, size(u), Val{N})
 end
@@ -75,7 +75,7 @@ function get_tmp(dc::FixedSizeDiffCache, u::Union{Number, AbstractArray})
     end
 end
 
-function get_tmp(dc::FixedSizeDiffCache, ::Type{T}) where T <: Number
+function get_tmp(dc::FixedSizeDiffCache, ::Type{T}) where {T <: Number}
     if promote_type(eltype(dc.du), T) <: eltype(dc.du)
         dc.du
     else
@@ -111,7 +111,7 @@ forward-mode automatic differentiation. Supports nested AD via keyword `levels`
 or specifying an array of chunk_sizes.
 """
 function DiffCache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u));
-    levels::Int = 1)
+        levels::Int = 1)
     DiffCache(u, size(u), N * ones(Int, levels))
 end
 DiffCache(u::AbstractArray, N::AbstractArray{<:Int}) = DiffCache(u, size(u), N)
@@ -164,7 +164,7 @@ function get_tmp(dc::DiffCache, u::Union{Number, AbstractArray})
     end
 end
 
-function get_tmp(dc::DiffCache, ::Type{T}) where T <: Number
+function get_tmp(dc::DiffCache, ::Type{T}) where {T <: Number}
     if promote_type(eltype(dc.du), T) <: eltype(dc.du)
         dc.du
     else
@@ -209,12 +209,18 @@ struct LazyBufferCache{F <: Function}
     LazyBufferCache(f::F = identity) where {F <: Function} = new{F}(Dict(), f) # start with empty dict
 end
 
+function similar_type(x::AbstractArray{T}, s::NTuple{N, Integer}) where {T, N}
+    # The compiler is smart enough to not allocate
+    # here for simple types like Array and SubArray
+    typeof(similar(x, ntuple(Returns(1), N)))
+end
+
 # override the [] method
 function Base.getindex(b::LazyBufferCache, u::T) where {T <: AbstractArray}
     s = b.sizemap(size(u)) # required buffer size
     get!(b.bufs, (T, s)) do
         similar(u, s) # buffer to allocate if it was not found in b.bufs
-    end::T  # declare type since b.bufs dictionary is untyped
+    end::similar_type(u, s) # declare type since b.bufs dictionary is untyped
 end
 
 # GeneralLazyBufferCache

test/general_lbc.jl

@@ -37,6 +37,9 @@ cache = LazyBufferCache()
 x = rand(1000)
 @inferred cache[x]
 @test 0 == @allocated cache[x]
+y = view(x, 1:900)
+@inferred cache[y]
+@test 0 == @allocated cache[y]
 
 cache = GeneralLazyBufferCache(T -> Vector{T}(undef, 1000))
 # GeneralLazyBufferCache is documented not to infer.