changed prefix sum to use an internal cache so it updates a list all at once using set_multiple

Author: adolgert

src/prefixsearch/binarytreeprefixsearch.jl

@@ -19,6 +19,8 @@ mutable struct BinaryTreePrefixSearch{T<:Real}
 	offset::Int64 # 2^(depth - 1). Index of first leaf and number of leaves.
 	cnt::Int64 # Number of leaves in use. Logical number of entries. cnt > 0.
 	initial_allocation::Int64
+	cache::Dict{Int64,T}
+	cached_cnt::Int64
 end
 
 
@@ -31,7 +33,7 @@ The optional hint, N, is the number of values to pre-allocate.
 function BinaryTreePrefixSearch{T}(N=32) where {T<:Real}
 	depth, offset, array_cnt = _btps_sizes(N)
 	b = zeros(T, array_cnt)
-	BinaryTreePrefixSearch{T}(b, depth, offset, 0, N)
+	BinaryTreePrefixSearch{T}(b, depth, offset, 0, N, Dict{Int64,T}(), 0)
 end
 
 
@@ -42,6 +44,8 @@ function Base.empty!(ps::BinaryTreePrefixSearch)
 	ps.depth = depth
 	ps.offset = offset
 	ps.cnt = 0
+	empty!(ps.cache)
+	ps.cached_cnt = 0
 end
 
 function Base.copy!(dst::BinaryTreePrefixSearch{T}, src::BinaryTreePrefixSearch{T}) where {T}
@@ -50,6 +54,8 @@ function Base.copy!(dst::BinaryTreePrefixSearch{T}, src::BinaryTreePrefixSearch{
 	dst.offset = src.offset
 	dst.cnt = src.cnt
 	dst.initial_allocation = src.initial_allocation
+	copy!(dst.cache, src.cache)
+	dst.cached_cnt = src.cached_cnt
 end
 
 
@@ -93,7 +99,7 @@ function Base.resize!(pst::BinaryTreePrefixSearch{T}, newcnt) where {T}
 end
 
 
-Base.length(ps::BinaryTreePrefixSearch) = ps.cnt
+Base.length(ps::BinaryTreePrefixSearch) = ps.cnt + ps.cached_cnt
 allocated(ps::BinaryTreePrefixSearch) = ps.offset
 
 
@@ -125,7 +131,15 @@ function choose(pst::BinaryTreePrefixSearch{T}, value) where {T}
 end
 
 
-Base.sum!(pst::BinaryTreePrefixSearch) = pst.array[1]
+# You have to call sum! before calling getindex, or it won't be updated. 
+function Base.sum!(pst::BinaryTreePrefixSearch)
+	if !isempty(pst.cache)
+		set_multiple!(pst, pairs(pst.cache))
+		empty!(pst.cache)
+		pst.cached_cnt = 0
+	end
+	pst.array[1]
+end
 
 
 """
@@ -161,16 +175,19 @@ end
 
 
 function Base.push!(pst::BinaryTreePrefixSearch{T}, value::T) where T
-	set_multiple!(pst, [(pst.cnt + 1, value)])
+	pst.cached_cnt += 1
+	pst.cache[pst.cnt + pst.cached_cnt] = value
 	return value
 end
 
 
 """
     setindex!(A, X, inds...)
+
+You have to call sum! before calling getindex.
 """
 function Base.setindex!(pst::BinaryTreePrefixSearch{T}, value::T, index) where T
-    set_multiple!(pst, [(index, value)])
+	pst.cache[index] = value
 end
 
 function Base.getindex(pst::BinaryTreePrefixSearch{T}, index) where {T}
@@ -200,8 +217,8 @@ Random.rand(rng::AbstractRNG, d::Random.SamplerTrivial{BinaryTreePrefixSearch{T}
 
 Base.haskey(md::BinaryTreePrefixSearch, clock) = false
 
-function Base.haskey(md::BinaryTreePrefixSearch{T}, clock::Int) where {T}
-	if 0 < clock ≤ length(md)
+function Base.haskey(pst::BinaryTreePrefixSearch{T}, clock::Int) where {T}
+	if 0 < clock ≤ length(pst)
         return getindex(pst, clock) > zero(T)
     else
         return false

test/test_direct.jl

@@ -39,6 +39,7 @@ end
 
     @test length(sampler) == 5
     @test length(keys(sampler)) == 5
+    next(sampler, 0.0, rng) # tells the sampler to rectify caches
     @test sampler[1] == 1 / 7.9
 
     @test haskey(sampler, 1)
@@ -48,6 +49,7 @@ end
     disable!(sampler, 1, 0.0)
 
     @test_throws KeyError sampler[1]
+    next(sampler, 0.0, rng)
     @test sampler[2] == 1 / 12.3
 
 end

test/test_prefixsearch.jl

@@ -73,18 +73,9 @@ end
     initial_allocation = 1
 	t = BinaryTreePrefixSearch{Int64}(initial_allocation)
     push!(t, 3)
-    @test t.cnt == 1
-    @test length(t) == 1
-    @test allocated(t) == 1
     push!(t, 4)
-    @test length(t) == 2
-    @test allocated(t) == 2
     push!(t, 2)
-    @test length(t) == 3
-    @test allocated(t) == 4
     push!(t, 3)
-    @test length(t) == 4
-    @test allocated(t) == 4
 	@test sum!(t) == 3+4+2+3
     @test t.array[1] == 3+4+2+3
     @test t.array[2] == 3+4
@@ -105,11 +96,13 @@ end
 	@test choose(t, 3)[1] == 2
 	@test choose(t, 3.7)[1] == 2
     t[1] = 4
+    sum!(t)
 	v = [(2, 1), (3.3, 1), (4.1, 2)]
 	for (guess, result) in v
 		@test choose(t, guess)[1] == result
 	end
     t[2] = 2
+    sum!(t)
 	v = [(2, 1), (3.3, 1), (5.1, 2)]
 	for (guess, result) in v
 		@test choose(t, guess)[1] == result
@@ -130,6 +123,7 @@ end
     @test choose(t, 5)[1] == 3
     t[2] = 2
     t[3] = 3
+    sum!(t)
     v = [(2, 1), (3, 2), (4.8, 2), (5.1, 3), (7.9, 3)]
     for (guess, result) in v
         @test choose(t, guess)[1] == result
@@ -149,6 +143,7 @@ end
     @test choose(t, 5.5)[1] == 3
     t[2] = 2.5
     t[3] = 3.5
+    sum!(t)
     v = [(2, 1),(3.5, 2),(4.8, 2),(6.1, 3),(7.9, 3)]
     for (guess, result) in v
         @test choose(t, guess)[1] == result
@@ -219,6 +214,7 @@ end
                 vals[idx] = rand(rng, 1:500)
                 btps[idx] = vals[idx]
             end
+            sum!(btps)
             same = all(btps[i] == vals[i] for i in eachindex(vals))
             @test same
             if !same