quadtree: optimize allocation in our maxHeap

benchmarks for move to our own maxHeap

benchmark                         old ns/op     new ns/op     delta
BenchmarkRandomKNearest10-12      4333          2094          -51.67%
BenchmarkRandomKNearest100-12     43245         15957         -63.10%

benchmark                         old allocs     new allocs     delta
BenchmarkRandomKNearest10-12      44             14             -68.18%
BenchmarkRandomKNearest100-12     369            104            -71.82%

benchmark                         old bytes     new bytes     delta
BenchmarkRandomKNearest10-12      1777          472           -73.44%
BenchmarkRandomKNearest100-12     15288         3432          -77.55%

quadtree/benchmarks_test.go

@@ -132,6 +132,23 @@ func BenchmarkRandomInBound1000Buf(b *testing.B) {
 	}
 }
 
+func BenchmarkRandomKNearest10(b *testing.B) {
+	r := rand.New(rand.NewSource(43))
+
+	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
+	for i := 0; i < 1000; i++ {
+		qt.Add(orb.Point{r.Float64(), r.Float64()})
+	}
+
+	buf := make([]orb.Pointer, 0, 10)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		qt.KNearest(buf[:0], orb.Point{r.Float64(), r.Float64()}, 10)
+	}
+}
+
 func BenchmarkRandomKNearest100(b *testing.B) {
 	r := rand.New(rand.NewSource(43))
 

quadtree/maxheap.go

@@ -13,27 +13,46 @@ type heapItem struct {
 	distance float64
 }
 
-func (h *maxHeap) Push(item *heapItem) {
-	*h = append(*h, item)
+func (h *maxHeap) Push(point orb.Pointer, distance float64) {
+	// Common usage is Push followed by a Pop if we have > k points.
+	// We're reusing the k+1 heapItem object to reduce memory allocations.
+	// First we manaully lengthen the slice,
+	// then we see if the last item has been allocated already.
+
+	prevLen := len(*h)
+	*h = (*h)[:prevLen+1]
+	if (*h)[prevLen] == nil {
+		(*h)[prevLen] = &heapItem{point: point, distance: distance}
+	} else {
+		(*h)[prevLen].point = point
+		(*h)[prevLen].distance = distance
+	}
 
 	i := len(*h) - 1
 	for i > 0 {
 		up := ((i + 1) >> 1) - 1
 		parent := (*h)[up]
 
-		if item.distance < parent.distance {
+		if distance < parent.distance {
 			// parent is further so we're done fixing up the heap.
 			break
 		}
 
 		// swap nodes
-		(*h)[i] = parent
-		(*h)[up] = item
+		// (*h)[i] = parent
+		(*h)[i].point = parent.point
+		(*h)[i].distance = parent.distance
+
+		// (*h)[up] = item
+		(*h)[up].point = point
+		(*h)[up].distance = distance
 
 		i = up
 	}
 }
 
+// Pop returns the "greatest" item in the list.
+// The returned item should not be saved across push/pop operations.
 func (h *maxHeap) Pop() *heapItem {
 	removed := (*h)[0]
 	lastItem := (*h)[len(*h)-1]

quadtree/maxheap_test.go

@@ -11,7 +11,7 @@ func TestMaxHeap(t *testing.T) {
 	for i := 1; i < 100; i++ {
 		h := make(maxHeap, 0, i)
 		for j := 0; j < i; j++ {
-			h.Push(&heapItem{distance: r.Float64()})
+			h.Push(nil, r.Float64())
 		}
 
 		current := h.Pop().distance

quadtree/quadtree.go

@@ -212,6 +212,7 @@ func (q *Quadtree) Matching(p orb.Point, f FilterFunc) orb.Pointer {
 // KNearest returns k closest Value/Pointer in the quadtree.
 // This function is thread safe. Multiple goroutines can read from a pre-created tree.
 // An optional buffer parameter is provided to allow for the reuse of result slice memory.
+// The points are returned in a sorted order, nearest first.
 // This function allows defining a maximum distance in order to reduce search iterations.
 func (q *Quadtree) KNearest(buf []orb.Pointer, p orb.Point, k int, maxDistance ...float64) []orb.Pointer {
 	return q.KNearestMatching(buf, p, k, nil, maxDistance...)
@@ -221,6 +222,7 @@ func (q *Quadtree) KNearest(buf []orb.Pointer, p orb.Point, k int, maxDistance .
 // the given filter function returns true. This function is thread safe.
 // Multiple goroutines can read from a pre-created tree. An optional buffer
 // parameter is provided to allow for the reuse of result slice memory.
+// The points are returned in a sorted order, nearest first.
 // This function allows defining a maximum distance in order to reduce search iterations.
 func (q *Quadtree) KNearestMatching(buf []orb.Pointer, p orb.Point, k int, f FilterFunc, maxDistance ...float64) []orb.Pointer {
 	if q.root == nil {
@@ -232,7 +234,7 @@ func (q *Quadtree) KNearestMatching(buf []orb.Pointer, p orb.Point, k int, f Fil
 		point:          p,
 		filter:         f,
 		k:              k,
-		maxHeap:        make(maxHeap, 0, k),
+		maxHeap:        make(maxHeap, 0, k+1),
 		closestBound:   &b,
 		maxDistSquared: math.MaxFloat64,
 	}
@@ -472,7 +474,7 @@ func (v *nearestVisitor) Visit(n *node) {
 
 	point := n.Value.Point()
 	if d := planar.DistanceSquared(point, v.point); d < v.maxDistSquared {
-		v.maxHeap.Push(&heapItem{point: n.Value, distance: d})
+		v.maxHeap.Push(n.Value, d)
 		if len(v.maxHeap) > v.k {
 
 			v.maxHeap.Pop()