Improve bounds checks in heap operations

Author: Drvi

src/heaps.jl

@@ -117,10 +117,10 @@ function nextreme(ord::Base.Ordering, n::Int, arr::AbstractVector{T}) where T
 
     rev = Base.ReverseOrdering(ord)
 
-    buffer = heapify(arr[1:n], rev)
+    buffer = heapify!(arr[1:n], rev)
 
-    for i = n + 1 : length(arr)
-        @inbounds xi = arr[i]
+    @inbounds for i = n + 1 : length(arr)
+        xi = arr[i]
         if Base.lt(rev, buffer[1], xi)
             buffer[1] = xi
             percolate_down!(buffer, 1, rev)

src/heaps/arrays_as_heaps.jl

@@ -13,33 +13,35 @@ heapleft(i::Integer) = 2i
 heapright(i::Integer) = 2i + 1
 heapparent(i::Integer) = div(i, 2)
 
-
 # Binary min-heap percolate down.
-function percolate_down!(xs::AbstractArray, i::Integer, x=xs[i], o::Ordering=Forward, len::Integer=length(xs))
+Base.@propagate_inbounds function percolate_down!(xs::AbstractArray, i::Integer, x, o::Ordering=Forward, len::Integer=length(xs))
+    @boundscheck checkbounds(xs, i)
+    @boundscheck checkbounds(xs, len)
+
     @inbounds while (l = heapleft(i)) <= len
         r = heapright(i)
         j = r > len || lt(o, xs[l], xs[r]) ? l : r
         lt(o, xs[j], x) || break
         xs[i] = xs[j]
         i = j
     end
-    xs[i] = x
+    @inbounds xs[i] = x
 end
-
-percolate_down!(xs::AbstractArray, i::Integer, o::Ordering, len::Integer=length(xs)) = percolate_down!(xs, i, xs[i], o, len)
+Base.@propagate_inbounds percolate_down!(xs::AbstractArray, i::Integer, o::Ordering, len::Integer=length(xs)) = percolate_down!(xs, i, xs[i], o, len)
 
 
 # Binary min-heap percolate up.
-function percolate_up!(xs::AbstractArray, i::Integer, x=xs[i], o::Ordering=Forward)
+Base.@propagate_inbounds function percolate_up!(xs::AbstractArray, i::Integer, x, o::Ordering=Forward)
+    @boundscheck checkbounds(xs, i)
+
     @inbounds while (j = heapparent(i)) >= 1
         lt(o, x, xs[j]) || break
         xs[i] = xs[j]
         i = j
     end
-    xs[i] = x
+    @inbounds xs[i] = x
 end
-
-@inline percolate_up!(xs::AbstractArray, i::Integer, o::Ordering) = percolate_up!(xs, i, xs[i], o)
+Base.@propagate_inbounds percolate_up!(xs::AbstractArray, i::Integer, o::Ordering) = percolate_up!(xs, i, xs[i], o)
 
 """
     heappop!(v, [ord])
@@ -48,10 +50,11 @@ Given a binary heap-ordered array, remove and return the lowest ordered element.
 For efficiency, this function does not check that the array is indeed heap-ordered.
 """
 function heappop!(xs::AbstractArray, o::Ordering=Forward)
+    Base.require_one_based_indexing(xs)
     x = xs[1]
     y = pop!(xs)
     if !isempty(xs)
-        percolate_down!(xs, 1, y, o)
+        @inbounds percolate_down!(xs, 1, y, o)
     end
     return x
 end
@@ -63,8 +66,9 @@ Given a binary heap-ordered array, push a new element `x`, preserving the heap p
 For efficiency, this function does not check that the array is indeed heap-ordered.
 """
 @inline function heappush!(xs::AbstractArray, x, o::Ordering=Forward)
+    Base.require_one_based_indexing(xs)
     push!(xs, x)
-    percolate_up!(xs, length(xs), o)
+    @inbounds percolate_up!(xs, length(xs), o)
     return xs
 end
 
@@ -76,8 +80,9 @@ end
 In-place [`heapify`](@ref).
 """
 @inline function heapify!(xs::AbstractArray, o::Ordering=Forward)
+    Base.require_one_based_indexing(xs)
     for i in heapparent(length(xs)):-1:1
-        percolate_down!(xs, i, o)
+        @inbounds percolate_down!(xs, i, o)
     end
     return xs
 end
@@ -129,6 +134,7 @@ false
 ```
 """
 function isheap(xs::AbstractArray, o::Ordering=Forward)
+    Base.require_one_based_indexing(xs)
     for i in 1:div(length(xs), 2)
         if lt(o, xs[heapleft(i)], xs[i]) ||
            (heapright(i) <= length(xs) && lt(o, xs[heapright(i)], xs[i]))

src/priorityqueue.jl

@@ -49,7 +49,7 @@ struct PriorityQueue{K,V,O<:Ordering} <: AbstractDict{K,V}
     function PriorityQueue{K,V,O}(o::O, itr) where {K,V,O<:Ordering}
         xs = Vector{Pair{K,V}}(undef, length(itr))
         index = Dict{K, Int}()
-        for (i, (k, v)) in enumerate(itr)
+        @inbounds for (i, (k, v)) in enumerate(itr)
             xs[i] = Pair{K,V}(k, v)
             if haskey(index, k)
                 throw(ArgumentError("PriorityQueue keys must be unique"))
@@ -60,7 +60,7 @@ struct PriorityQueue{K,V,O<:Ordering} <: AbstractDict{K,V}
 
         # heapify
         for i in heapparent(length(pq.xs)):-1:1
-            percolate_down!(pq, i)
+            @inbounds percolate_down!(pq, i)
         end
 
         return pq
@@ -167,8 +167,10 @@ priority queue.
 """
 Base.first(pq::PriorityQueue) = first(pq.xs)
 
-function percolate_down!(pq::PriorityQueue, i::Integer)
-    x = pq.xs[i]
+Base.@propagate_inbounds function percolate_down!(pq::PriorityQueue, i::Integer)
+    @boundscheck checkbounds(pq.xs, i)
+
+    @inbounds x = pq.xs[i]
     @inbounds while (l = heapleft(i)) <= length(pq)
         r = heapright(i)
         j = r > length(pq) || lt(pq.o, pq.xs[l].second, pq.xs[r].second) ? l : r
@@ -182,12 +184,14 @@ function percolate_down!(pq::PriorityQueue, i::Integer)
         end
     end
     pq.index[x.first] = i
-    pq.xs[i] = x
+    @inbounds pq.xs[i] = x
 end
 
 
-function percolate_up!(pq::PriorityQueue, i::Integer)
-    x = pq.xs[i]
+Base.@propagate_inbounds function percolate_up!(pq::PriorityQueue, i::Integer)
+    @boundscheck checkbounds(pq.xs, i)
+
+    @inbounds x = pq.xs[i]
     @inbounds while i > 1
         j = heapparent(i)
         xj = pq.xs[j]
@@ -200,7 +204,7 @@ function percolate_up!(pq::PriorityQueue, i::Integer)
         end
     end
     pq.index[x.first] = i
-    pq.xs[i] = x
+    @inbounds pq.xs[i] = x
 end
 
 # Equivalent to percolate_up! with an element having lower priority than any other
@@ -236,8 +240,8 @@ end
 # Change the priority of an existing element, or enqueue it if it isn't present.
 function Base.setindex!(pq::PriorityQueue{K, V}, value, key) where {K,V}
     i = get(pq.index, key, 0)
-    if i != 0
-        @inbounds oldvalue = pq.xs[i].second
+    @inbounds if i != 0
+        oldvalue = pq.xs[i].second
         pq.xs[i] = Pair{K,V}(key, value)
         if lt(pq.o, oldvalue, value)
             percolate_down!(pq, i)
@@ -255,7 +259,7 @@ end
 
 Insert the a key `k` into a priority queue `pq` with priority `v`.
 
-# Examples 
+# Examples
 
 ```jldoctest
 julia> a = PriorityQueue("a" => 1, "b" => 2, "c" => 3, "e" => 5)
@@ -317,7 +321,7 @@ function Base.popfirst!(pq::PriorityQueue)
     if !isempty(pq)
         @inbounds pq.xs[1] = y
         pq.index[y.first] = 1
-        percolate_down!(pq, 1)
+        @inbounds percolate_down!(pq, 1)
     end
     delete!(pq.index, x.first)
     return x
@@ -354,7 +358,7 @@ function Base.delete!(pq::PriorityQueue, key)
 end
 
 """
-    empty!(pq::PriorityQueue)  
+    empty!(pq::PriorityQueue)
 
 Reset priority queue `pq`.
 """