quadtree: fix cleanup of nodes during removal

quadtree/quadtree.go

@@ -64,6 +64,9 @@ func (q *Quadtree) Add(p orb.Pointer) error {
 			Value: p,
 		}
 		return nil
+	} else if q.root.Value == nil {
+		q.root.Value = p
+		return nil
 	}
 
 	q.add(q.root, p, p.Point(),
@@ -98,6 +101,9 @@ func (q *Quadtree) add(n *node, p orb.Pointer, point orb.Point, left, right, bot
 	if n.Children[i] == nil {
 		n.Children[i] = &node{Value: p}
 		return
+	} else if n.Children[i].Value == nil {
+		n.Children[i].Value = p
+		return
 	}
 
 	// proceed down to the child to see if it's a leaf yet and we can add the pointer there.
@@ -137,40 +143,49 @@ func (q *Quadtree) Remove(p orb.Pointer, eq FilterFunc) bool {
 		return false
 	}
 
+	v.closest.Value = nil
+
+	// if v.closest is NOT a leaf node, values will be shuffled up into this node.
+	// if v.closest IS a leaf node, the call is a no-op but we can't delete
+	// the now empty node because we don't know the parent here.
+	//
+	// Future adds will reuse this node if applicable.
+	// Removing v.closest parent will cause this node to be removed,
+	// but the parent will be a leaf with a nil value.
 	removeNode(v.closest)
 	return true
 }
 
 // removeNode is the recursive fixing up of the tree when we remove a node.
-func removeNode(n *node) {
-	var i int
-	for {
-		i = -1
-		if n.Children[0] != nil {
-			i = 0
-		} else if n.Children[1] != nil {
-			i = 1
-		} else if n.Children[2] != nil {
-			i = 2
-		} else if n.Children[3] != nil {
-			i = 3
-		}
+// It will pull up a child value into it's place. It will try to remove leaf nodes
+// that are now empty, since their values got pulled up.
+func removeNode(n *node) bool {
+	i := -1
+	if n.Children[0] != nil {
+		i = 0
+	} else if n.Children[1] != nil {
+		i = 1
+	} else if n.Children[2] != nil {
+		i = 2
+	} else if n.Children[3] != nil {
+		i = 3
+	}
 
-		if i == -1 {
-			n.Value = nil
-			return
-		}
+	if i == -1 {
+		// all children are nil, can remove.
+		// n.value ==  nil because it "pulled up" (or removed) by the caller.
+		return true
+	}
 
-		if n.Children[i].Value == nil {
-			n.Children[i] = nil
-			continue
-		}
+	n.Value = n.Children[i].Value
+	n.Children[i].Value = nil
 
-		break
+	removeThisChild := removeNode(n.Children[i])
+	if removeThisChild {
+		n.Children[i] = nil
 	}
 
-	n.Value = n.Children[i].Value
-	removeNode(n.Children[i])
+	return false
 }
 
 // Find returns the closest Value/Pointer in the quadtree.

quadtree/quadtree_test.go

@@ -1,10 +1,12 @@
 package quadtree
 
 import (
+	"fmt"
 	"math/rand"
 	"reflect"
 	"sort"
 	"testing"
+	"time"
 
 	"github.com/paulmach/orb"
 	"github.com/paulmach/orb/planar"
@@ -56,6 +58,232 @@ func TestQuadtreeRemove(t *testing.T) {
 	}
 }
 
+type PExtra struct {
+	p  orb.Point
+	id string
+}
+
+func (p *PExtra) Point() orb.Point {
+	return p.p
+}
+
+func (p *PExtra) String() string {
+	return fmt.Sprintf("%v: %v", p.id, p.p)
+}
+
+func TestQuadtreeRemoveAndAdd_inOrder(t *testing.T) {
+	seed := time.Now().UnixNano()
+	t.Logf("seed: %v", seed)
+	r := rand.New(rand.NewSource(seed))
+
+	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
+	p1 := &PExtra{p: orb.Point{r.Float64(), r.Float64()}, id: "1"}
+	p2 := &PExtra{p: orb.Point{p1.p[0], p1.p[1]}, id: "2"}
+	p3 := &PExtra{p: orb.Point{p1.p[0], p1.p[1]}, id: "3"}
+
+	qt.Add(p1)
+	qt.Add(p2)
+	qt.Add(p3)
+
+	// rm 3
+	found := qt.Remove(p3, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p3.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	// leaf node doesn't actually get removed
+	if c := countNodes(qt.root); c != 3 {
+		t.Errorf("incorrect number of nodes: %v != 3", c)
+	}
+
+	// 3 again
+	found = qt.Remove(p3, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p3.id
+	})
+	if found {
+		t.Errorf("should not find already removed node")
+	}
+
+	// rm 2
+	found = qt.Remove(p2, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p2.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	if c := countNodes(qt.root); c != 2 {
+		t.Errorf("incorrect number of nodes: %v != 2", c)
+	}
+
+	// rm 1
+	found = qt.Remove(p1, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p1.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	if c := countNodes(qt.root); c != 1 {
+		t.Errorf("incorrect number of nodes: %v != 1", c)
+	}
+}
+
+func TestQuadtreeRemoveAndAdd_sameLoc(t *testing.T) {
+	seed := time.Now().UnixNano()
+	t.Logf("seed: %v", seed)
+	r := rand.New(rand.NewSource(seed))
+
+	qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}})
+	p1 := &PExtra{p: orb.Point{r.Float64(), r.Float64()}, id: "1"}
+	p2 := &PExtra{p: orb.Point{p1.p[0], p1.p[1]}, id: "2"}
+	p3 := &PExtra{p: orb.Point{p1.p[0], p1.p[1]}, id: "3"}
+	p4 := &PExtra{p: orb.Point{p1.p[0], p1.p[1]}, id: "4"}
+	p5 := &PExtra{p: orb.Point{p1.p[0], p1.p[1]}, id: "5"}
+
+	qt.Add(p1)
+	qt.Add(p2)
+	qt.Add(p3)
+
+	// remove middle point
+	found := qt.Remove(p2, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p2.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	if c := countNodes(qt.root); c != 2 {
+		t.Errorf("incorrect number of nodes: %v != 2", c)
+	}
+
+	// remove first point
+	found = qt.Remove(p1, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p1.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	if c := countNodes(qt.root); c != 1 {
+		t.Errorf("incorrect number of nodes: %v != 1", c)
+	}
+
+	// add a 4th point
+	qt.Add(p4)
+
+	// remove third point
+	found = qt.Remove(p3, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p3.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	if c := countNodes(qt.root); c != 1 {
+		t.Errorf("incorrect number of nodes: %v != 1", c)
+	}
+
+	// add a 5th point
+	qt.Add(p5)
+
+	// remove the 5th point
+	found = qt.Remove(p5, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p5.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	// 5 is a tail point, so its not does not actually get removed
+	if c := countNodes(qt.root); c != 2 {
+		t.Errorf("incorrect number of nodes: %v != 2", c)
+	}
+
+	// add a 3th point again
+	qt.Add(p3)
+
+	// should reuse the tail point left by p5
+	if c := countNodes(qt.root); c != 2 {
+		t.Errorf("incorrect number of nodes: %v != 2", c)
+	}
+
+	// remove p4/root
+	found = qt.Remove(p4, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p4.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	if c := countNodes(qt.root); c != 1 {
+		t.Errorf("incorrect number of nodes: %v != 1", c)
+	}
+
+	// remove p3/root
+	found = qt.Remove(p3, func(p orb.Pointer) bool {
+		return p.(*PExtra).id == p3.id
+	})
+	if !found {
+		t.Error("didn't find/remove point")
+	}
+
+	// just the root, can't remove it
+	if c := countNodes(qt.root); c != 1 {
+		t.Errorf("incorrect number of nodes: %v != 1", c)
+	}
+
+	// add back a point to be put in the root
+	qt.Add(p3)
+
+	if c := countNodes(qt.root); c != 1 {
+		t.Errorf("incorrect number of nodes: %v != 1", c)
+	}
+}
+
+func TestQuadtreeRemoveAndAdd_random(t *testing.T) {
+	seed := time.Now().UnixNano()
+	t.Logf("seed: %v", seed)
+	r := rand.New(rand.NewSource(seed))
+
+	const runs = 10
+	const perRun = 300 // add 300, remove 300/2
+
+	bounds := orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{3000, 3000}}
+	qt := New(bounds)
+	points := make([]*PExtra, 0, 3000)
+	id := 0
+
+	for i := 0; i < runs; i++ {
+		for j := 0; j < perRun; j++ {
+			x := r.Int63n(30)
+			y := r.Int63n(30)
+			id++
+			p := &PExtra{p: orb.Point{float64(x), float64(y)}, id: fmt.Sprintf("%d", id)}
+			qt.Add(p)
+			points = append(points, p)
+
+		}
+
+		for j := 0; j < perRun/2; j++ {
+			k := r.Int() % len(points)
+			remP := points[k]
+			points = append(points[:k], points[k+1:]...)
+			qt.Remove(remP, func(p orb.Pointer) bool {
+				return p.(*PExtra).id == remP.id
+			})
+		}
+	}
+
+	left := len(qt.InBound(nil, bounds))
+	expected := runs * perRun / 2
+	if left != expected {
+		t.Errorf("incorrect number of points in tree: %d != %d", left, expected)
+	}
+}
+
 func TestQuadtreeFind(t *testing.T) {
 	points := orb.MultiPoint{}
 	dim := 17
@@ -426,3 +654,17 @@ func TestQuadtreeInBound_Random(t *testing.T) {
 		}
 	}
 }
+
+func countNodes(n *node) int {
+	if n == nil {
+		return 0
+	}
+
+	c := 1
+	c += countNodes(n.Children[0])
+	c += countNodes(n.Children[1])
+	c += countNodes(n.Children[2])
+	c += countNodes(n.Children[3])
+
+	return c
+}