diff --git a/internal/collection/collection.go b/internal/collection/collection.go index 0833e574..bfb46ee9 100644 --- a/internal/collection/collection.go +++ b/internal/collection/collection.go @@ -1,14 +1,15 @@ package collection import ( + "reflect" "unsafe" - "github.com/tidwall/boxtree/d2" "github.com/tidwall/btree" "github.com/tidwall/geojson" "github.com/tidwall/geojson/geo" "github.com/tidwall/geojson/geometry" "github.com/tidwall/tile38/internal/collection/ptrbtree" + "github.com/tidwall/tile38/internal/collection/ptrrtree" ) // Cursor allows for quickly paging through Scan, Within, Intersects, and Nearby @@ -18,11 +19,32 @@ type Cursor interface { } type itemT struct { - id string obj geojson.Object + _ uint32 + idLen uint32 + idData unsafe.Pointer fields []float64 } +func (item *itemT) id() string { + return *(*string)((unsafe.Pointer)(&reflect.StringHeader{ + Data: uintptr(unsafe.Pointer(item.idData)), + Len: int(item.idLen), + })) +} + +func newItem(id string, obj geojson.Object) *itemT { + item := new(itemT) + item.obj = obj + item.idLen = uint32(len(id)) + if len(id) > 0 { + idData := make([]byte, len(id)) + copy(idData, id) + item.idData = unsafe.Pointer(&idData[0]) + } + return item +} + func (item *itemT) weightAndPoints() (weight, points int) { if objIsSpatial(item.obj) { points = item.obj.NumPoints() @@ -30,7 +52,7 @@ func (item *itemT) weightAndPoints() (weight, points int) { } else { weight = len(item.obj.String()) } - weight += len(item.fields)*8 + len(item.id) + weight += len(item.fields)*8 + len(item.id()) return weight, points } @@ -44,14 +66,14 @@ func (item *itemT) Less(other btree.Item, ctx interface{}) bool { return false } // the values match so we'll compare IDs, which are always unique. - return item.id < other.(*itemT).id + return item.id() < other.(*itemT).id() } // Collection represents a collection of geojson objects. type Collection struct { - items ptrbtree.BTree // items sorted by keys - index d2.BoxTree // items geospatially indexed - values *btree.BTree // items sorted by value+key + items ptrbtree.BTree // items sorted by keys + index ptrrtree.BoxTree // items geospatially indexed + values *btree.BTree // items sorted by value+key fieldMap map[string]int weight int points int @@ -110,7 +132,7 @@ func (c *Collection) indexDelete(item *itemT) { c.index.Delete( []float64{rect.Min.X, rect.Min.Y}, []float64{rect.Max.X, rect.Max.Y}, - item) + unsafe.Pointer(item)) } } @@ -120,7 +142,7 @@ func (c *Collection) indexInsert(item *itemT) { c.index.Insert( []float64{rect.Min.X, rect.Min.Y}, []float64{rect.Max.X, rect.Max.Y}, - item) + unsafe.Pointer(item)) } } @@ -160,7 +182,7 @@ func (c *Collection) Set( ) ( oldObject geojson.Object, oldFields []float64, newFields []float64, ) { - newItem := &itemT{id: id, obj: obj} + newItem := newItem(id, obj) // add the new item to main btree and remove the old one if needed oldItemV, ok := c.items.Set(unsafe.Pointer(newItem)) @@ -336,7 +358,7 @@ func (c *Collection) Scan(desc bool, cursor Cursor, cursor.Step(1) } iitm := (*itemT)(ptr) - keepon = iterator(iitm.id, iitm.obj, iitm.fields) + keepon = iterator(iitm.id(), iitm.obj, iitm.fields) return keepon } if desc { @@ -368,15 +390,15 @@ func (c *Collection) ScanRange(start, end string, desc bool, cursor Cursor, } iitm := (*itemT)(ptr) if !desc { - if iitm.id >= end { + if iitm.id() >= end { return false } } else { - if iitm.id <= end { + if iitm.id() <= end { return false } } - keepon = iterator(iitm.id, iitm.obj, iitm.fields) + keepon = iterator(iitm.id(), iitm.obj, iitm.fields) return keepon } @@ -408,7 +430,7 @@ func (c *Collection) SearchValues(desc bool, cursor Cursor, cursor.Step(1) } iitm := item.(*itemT) - keepon = iterator(iitm.id, iitm.obj, iitm.fields) + keepon = iterator(iitm.id(), iitm.obj, iitm.fields) return keepon } if desc { @@ -440,15 +462,17 @@ func (c *Collection) SearchValuesRange(start, end string, desc bool, cursor.Step(1) } iitm := item.(*itemT) - keepon = iterator(iitm.id, iitm.obj, iitm.fields) + keepon = iterator(iitm.id(), iitm.obj, iitm.fields) return keepon } if desc { - c.values.DescendRange(&itemT{obj: String(start)}, - &itemT{obj: String(end)}, iter) + c.values.DescendRange( + newItem("", String(start)), newItem("", String(end)), iter, + ) } else { - c.values.AscendRange(&itemT{obj: String(start)}, - &itemT{obj: String(end)}, iter) + c.values.AscendRange( + newItem("", String(start)), newItem("", String(end)), iter, + ) } return keepon } @@ -474,7 +498,7 @@ func (c *Collection) ScanGreaterOrEqual(id string, desc bool, cursor.Step(1) } iitm := (*itemT)(ptr) - keepon = iterator(iitm.id, iitm.obj, iitm.fields) + keepon = iterator(iitm.id(), iitm.obj, iitm.fields) return keepon } if desc { @@ -493,9 +517,9 @@ func (c *Collection) geoSearch( c.index.Search( []float64{rect.Min.X, rect.Min.Y}, []float64{rect.Max.X, rect.Max.Y}, - func(_, _ []float64, itemv interface{}) bool { - item := itemv.(*itemT) - alive = iter(item.id, item.obj, item.fields) + func(_, _ []float64, itemv unsafe.Pointer) bool { + item := (*itemT)(itemv) + alive = iter(item.id(), item.obj, item.fields) return alive }, ) @@ -688,7 +712,7 @@ func (c *Collection) Nearby( c.index.Search( []float64{minLon, minLat}, []float64{maxLon, maxLat}, - func(_, _ []float64, itemv interface{}) bool { + func(_, _ []float64, itemv unsafe.Pointer) bool { exists = true return false }, @@ -711,7 +735,7 @@ func (c *Collection) Nearby( c.index.Nearby( []float64{center.X, center.Y}, []float64{center.X, center.Y}, - func(_, _ []float64, itemv interface{}) bool { + func(_, _ []float64, itemv unsafe.Pointer) bool { count++ if count <= offset { return true @@ -719,8 +743,8 @@ func (c *Collection) Nearby( if cursor != nil { cursor.Step(1) } - item := itemv.(*itemT) - alive = iter(item.id, item.obj, item.fields) + item := (*itemT)(itemv) + alive = iter(item.id(), item.obj, item.fields) return alive }, ) diff --git a/internal/collection/ptrbtree/btree.go b/internal/collection/ptrbtree/btree.go index 9f4ef6bf..9c09f9b3 100644 --- a/internal/collection/ptrbtree/btree.go +++ b/internal/collection/ptrbtree/btree.go @@ -1,20 +1,36 @@ package ptrbtree -import "unsafe" +import ( + "reflect" + "unsafe" +) const maxItems = 31 // use an odd number const minItems = maxItems * 40 / 100 -type item struct{ ptr unsafe.Pointer } -type keyedItem struct{ key string } +type btreeItem struct { + ptr unsafe.Pointer +} -func (v item) key() string { - return (*keyedItem)(v.ptr).key +// keyedItem must match layout of ../collection/itemT, otherwise +// there's a risk for memory corruption. +type keyedItem struct { + _ interface{} + _ uint32 + keyLen uint32 + data unsafe.Pointer +} + +func (v btreeItem) key() string { + return *(*string)((unsafe.Pointer)(&reflect.StringHeader{ + Data: uintptr(unsafe.Pointer((*keyedItem)(v.ptr).data)), + Len: int((*keyedItem)(v.ptr).keyLen), + })) } type node struct { numItems int - items [maxItems]item + items [maxItems]btreeItem children [maxItems + 1]*node } @@ -44,7 +60,7 @@ func (n *node) find(key string) (index int, found bool) { // Set or replace a value for a key func (tr *BTree) Set(ptr unsafe.Pointer) (prev unsafe.Pointer, replaced bool) { - newItem := item{ptr} + newItem := btreeItem{ptr} if tr.root == nil { tr.root = new(node) tr.root.items[0] = newItem @@ -70,7 +86,7 @@ func (tr *BTree) Set(ptr unsafe.Pointer) (prev unsafe.Pointer, replaced bool) { return } -func (n *node) split(height int) (right *node, median item) { +func (n *node) split(height int) (right *node, median btreeItem) { right = new(node) median = n.items[maxItems/2] copy(right.items[:maxItems/2], n.items[maxItems/2+1:]) @@ -84,13 +100,13 @@ func (n *node) split(height int) (right *node, median item) { } } for i := maxItems / 2; i < maxItems; i++ { - n.items[i] = item{} + n.items[i] = btreeItem{} } n.numItems = maxItems / 2 return } -func (n *node) set(newItem item, height int) (prev unsafe.Pointer, replaced bool) { +func (n *node) set(newItem btreeItem, height int) (prev unsafe.Pointer, replaced bool) { i, found := n.find(newItem.key()) if found { prev = n.items[i].ptr @@ -176,7 +192,7 @@ func (tr *BTree) Delete(key string) (prev unsafe.Pointer, deleted bool) { if tr.root == nil { return } - var prevItem item + var prevItem btreeItem prevItem, deleted = tr.root.delete(false, key, tr.height) if !deleted { return @@ -195,7 +211,7 @@ func (tr *BTree) Delete(key string) (prev unsafe.Pointer, deleted bool) { } func (n *node) delete(max bool, key string, height int) ( - prev item, deleted bool, + prev btreeItem, deleted bool, ) { i, found := 0, false if max { @@ -208,12 +224,12 @@ func (n *node) delete(max bool, key string, height int) ( prev = n.items[i] // found the items at the leaf, remove it and return. copy(n.items[i:], n.items[i+1:n.numItems]) - n.items[n.numItems-1] = item{} + n.items[n.numItems-1] = btreeItem{} n.children[n.numItems] = nil n.numItems-- return prev, true } - return item{}, false + return btreeItem{}, false } if found { @@ -237,7 +253,7 @@ func (n *node) delete(max bool, key string, height int) ( i-- } if n.children[i].numItems+n.children[i+1].numItems+1 < maxItems { - // merge left + item + right + // merge left + btreeItem + right n.children[i].items[n.children[i].numItems] = n.items[i] copy(n.children[i].items[n.children[i].numItems+1:], n.children[i+1].items[:n.children[i+1].numItems]) @@ -248,7 +264,7 @@ func (n *node) delete(max bool, key string, height int) ( n.children[i].numItems += n.children[i+1].numItems + 1 copy(n.items[i:], n.items[i+1:n.numItems]) copy(n.children[i+1:], n.children[i+2:n.numItems+1]) - n.items[n.numItems] = item{} + n.items[n.numItems] = btreeItem{} n.children[n.numItems+1] = nil n.numItems-- } else if n.children[i].numItems > n.children[i+1].numItems { @@ -266,7 +282,7 @@ func (n *node) delete(max bool, key string, height int) ( } n.children[i+1].numItems++ n.items[i] = n.children[i].items[n.children[i].numItems-1] - n.children[i].items[n.children[i].numItems-1] = item{} + n.children[i].items[n.children[i].numItems-1] = btreeItem{} if height > 1 { n.children[i].children[n.children[i].numItems] = nil } diff --git a/internal/collection/ptrrtree/rtree.go b/internal/collection/ptrrtree/rtree.go new file mode 100644 index 00000000..cd86e05f --- /dev/null +++ b/internal/collection/ptrrtree/rtree.go @@ -0,0 +1,712 @@ +package ptrrtree + +import "unsafe" + +const dims = 2 + +const ( + maxEntries = 16 + minEntries = maxEntries * 40 / 100 +) + +type box struct { + data unsafe.Pointer + min, max [dims]float64 +} + +type node struct { + count int + boxes [maxEntries + 1]box +} + +// BoxTree ... +type BoxTree struct { + height int + root box + count int + reinsert []box +} + +func (r *box) expand(b *box) { + for i := 0; i < dims; i++ { + if b.min[i] < r.min[i] { + r.min[i] = b.min[i] + } + if b.max[i] > r.max[i] { + r.max[i] = b.max[i] + } + } +} + +func (r *box) area() float64 { + area := r.max[0] - r.min[0] + for i := 1; i < dims; i++ { + area *= r.max[i] - r.min[i] + } + return area +} + +func (r *box) overlapArea(b *box) float64 { + area := 1.0 + for i := 0; i < dims; i++ { + var max, min float64 + if r.max[i] < b.max[i] { + max = r.max[i] + } else { + max = b.max[i] + } + if r.min[i] > b.min[i] { + min = r.min[i] + } else { + min = b.min[i] + } + if max > min { + area *= max - min + } else { + return 0 + } + } + return area +} + +func (r *box) enlargedArea(b *box) float64 { + area := 1.0 + for i := 0; i < len(r.min); i++ { + if b.max[i] > r.max[i] { + if b.min[i] < r.min[i] { + area *= b.max[i] - b.min[i] + } else { + area *= b.max[i] - r.min[i] + } + } else { + if b.min[i] < r.min[i] { + area *= r.max[i] - b.min[i] + } else { + area *= r.max[i] - r.min[i] + } + } + } + return area +} + +// Insert inserts an item into the RTree +func (tr *BoxTree) Insert(min, max []float64, value unsafe.Pointer) { + var item box + fit(min, max, value, &item) + tr.insert(&item) +} + +func (tr *BoxTree) insert(item *box) { + if tr.root.data == nil { + fit(item.min[:], item.max[:], unsafe.Pointer(new(node)), &tr.root) + } + grown := tr.root.insert(item, tr.height) + if grown { + tr.root.expand(item) + } + if (*node)(tr.root.data).count == maxEntries+1 { + newRoot := new(node) + tr.root.splitLargestAxisEdgeSnap(&newRoot.boxes[1]) + newRoot.boxes[0] = tr.root + newRoot.count = 2 + tr.root.data = unsafe.Pointer(newRoot) + tr.root.recalc() + tr.height++ + } + tr.count++ +} + +func (r *box) chooseLeastEnlargement(b *box) int { + j, jenlargement, jarea := -1, 0.0, 0.0 + n := (*node)(r.data) + for i := 0; i < n.count; i++ { + var area float64 + if false { + area = n.boxes[i].area() + } else { + // force inline + area = n.boxes[i].max[0] - n.boxes[i].min[0] + for j := 1; j < dims; j++ { + area *= n.boxes[i].max[j] - n.boxes[i].min[j] + } + } + var enlargement float64 + if false { + enlargement = n.boxes[i].enlargedArea(b) - area + } else { + // force inline + enlargedArea := 1.0 + for j := 0; j < len(n.boxes[i].min); j++ { + if b.max[j] > n.boxes[i].max[j] { + if b.min[j] < n.boxes[i].min[j] { + enlargedArea *= b.max[j] - b.min[j] + } else { + enlargedArea *= b.max[j] - n.boxes[i].min[j] + } + } else { + if b.min[j] < n.boxes[i].min[j] { + enlargedArea *= n.boxes[i].max[j] - b.min[j] + } else { + enlargedArea *= n.boxes[i].max[j] - n.boxes[i].min[j] + } + } + } + enlargement = enlargedArea - area + } + + if j == -1 || enlargement < jenlargement { + j, jenlargement, jarea = i, enlargement, area + } else if enlargement == jenlargement { + if area < jarea { + j, jenlargement, jarea = i, enlargement, area + } + } + } + return j +} + +func (r *box) recalc() { + n := (*node)(r.data) + r.min = n.boxes[0].min + r.max = n.boxes[0].max + for i := 1; i < n.count; i++ { + r.expand(&n.boxes[i]) + } +} + +// contains return struct when b is fully contained inside of n +func (r *box) contains(b *box) bool { + for i := 0; i < dims; i++ { + if b.min[i] < r.min[i] || b.max[i] > r.max[i] { + return false + } + } + return true +} + +func (r *box) largestAxis() (axis int, size float64) { + j, jsz := 0, 0.0 + for i := 0; i < dims; i++ { + sz := r.max[i] - r.min[i] + if i == 0 || sz > jsz { + j, jsz = i, sz + } + } + return j, jsz +} + +func (r *box) splitLargestAxisEdgeSnap(right *box) { + axis, _ := r.largestAxis() + left := r + leftNode := (*node)(left.data) + rightNode := new(node) + right.data = unsafe.Pointer(rightNode) + + var equals []box + for i := 0; i < leftNode.count; i++ { + minDist := leftNode.boxes[i].min[axis] - left.min[axis] + maxDist := left.max[axis] - leftNode.boxes[i].max[axis] + if minDist < maxDist { + // stay left + } else { + if minDist > maxDist { + // move to right + rightNode.boxes[rightNode.count] = leftNode.boxes[i] + rightNode.count++ + } else { + // move to equals, at the end of the left array + equals = append(equals, leftNode.boxes[i]) + } + leftNode.boxes[i] = leftNode.boxes[leftNode.count-1] + leftNode.boxes[leftNode.count-1].data = nil + leftNode.count-- + i-- + } + } + for _, b := range equals { + if leftNode.count < rightNode.count { + leftNode.boxes[leftNode.count] = b + leftNode.count++ + } else { + rightNode.boxes[rightNode.count] = b + rightNode.count++ + } + } + left.recalc() + right.recalc() +} + +func (r *box) insert(item *box, height int) (grown bool) { + n := (*node)(r.data) + if height == 0 { + n.boxes[n.count] = *item + n.count++ + grown = !r.contains(item) + return grown + } + // choose subtree + index := r.chooseLeastEnlargement(item) + child := &n.boxes[index] + grown = child.insert(item, height-1) + if grown { + child.expand(item) + grown = !r.contains(item) + } + if (*node)(child.data).count == maxEntries+1 { + child.splitLargestAxisEdgeSnap(&n.boxes[n.count]) + n.count++ + } + return grown +} + +// fit an external item into a box type +func fit(min, max []float64, value unsafe.Pointer, target *box) { + if max == nil { + max = min + } + if len(min) != len(max) { + panic("min/max dimension mismatch") + } + if len(min) != dims { + panic("invalid number of dimensions") + } + for i := 0; i < dims; i++ { + target.min[i] = min[i] + target.max[i] = max[i] + } + target.data = value +} + +type overlapsResult int + +const ( + not overlapsResult = iota + intersects + contains +) + +// overlaps detects if r insersects or contains b. +// return not, intersects, contains +func (r *box) overlaps(b *box) overlapsResult { + for i := 0; i < dims; i++ { + if b.min[i] > r.max[i] || b.max[i] < r.min[i] { + return not + } + if r.min[i] > b.min[i] || b.max[i] > r.max[i] { + i++ + for ; i < dims; i++ { + if b.min[i] > r.max[i] || b.max[i] < r.min[i] { + return not + } + } + return intersects + } + } + return contains +} + +// contains return struct when b is fully contained inside of n +func (r *box) intersects(b *box) bool { + for i := 0; i < dims; i++ { + if b.min[i] > r.max[i] || b.max[i] < r.min[i] { + return false + } + } + return true +} + +func (r *box) search( + target *box, height int, + iter func(min, max []float64, value unsafe.Pointer) bool, +) bool { + n := (*node)(r.data) + if height == 0 { + for i := 0; i < n.count; i++ { + if target.intersects(&n.boxes[i]) { + if !iter(n.boxes[i].min[:], n.boxes[i].max[:], + n.boxes[i].data) { + return false + } + } + } + } else { + for i := 0; i < n.count; i++ { + switch target.overlaps(&n.boxes[i]) { + case intersects: + if !n.boxes[i].search(target, height-1, iter) { + return false + } + case contains: + if !n.boxes[i].scan(height-1, iter) { + return false + } + } + } + } + return true +} + +func (tr *BoxTree) search( + target *box, + iter func(min, max []float64, value unsafe.Pointer) bool, +) { + if tr.root.data == nil { + return + } + res := target.overlaps(&tr.root) + if res == intersects { + tr.root.search(target, tr.height, iter) + } else if res == contains { + tr.root.scan(tr.height, iter) + } +} + +// Search ... +func (tr *BoxTree) Search(min, max []float64, + iter func(min, max []float64, value unsafe.Pointer) bool, +) { + var target box + fit(min, max, nil, &target) + tr.search(&target, iter) +} + +const ( + // Continue to first child box and/or next sibling. + Continue = iota + // Ignore child boxes but continue to next sibling. + Ignore + // Stop iterating + Stop +) + +// Traverse iterates through all items and container boxes in tree. +func (tr *BoxTree) Traverse( + iter func(min, max []float64, height, level int, value unsafe.Pointer) int, +) { + if tr.root.data == nil { + return + } + if iter(tr.root.min[:], tr.root.max[:], tr.height+1, 0, nil) == Continue { + tr.root.traverse(tr.height, 1, iter) + } +} + +func (r *box) traverse( + height, level int, + iter func(min, max []float64, height, level int, value unsafe.Pointer) int, +) int { + n := (*node)(r.data) + if height == 0 { + for i := 0; i < n.count; i++ { + action := iter(n.boxes[i].min[:], n.boxes[i].max[:], height, level, + n.boxes[i].data) + if action == Stop { + return Stop + } + } + } else { + for i := 0; i < n.count; i++ { + switch iter(n.boxes[i].min[:], n.boxes[i].max[:], height, level, + n.boxes[i].data) { + case Ignore: + case Continue: + if n.boxes[i].traverse(height-1, level+1, iter) == Stop { + return Stop + } + case Stop: + return Stop + } + } + } + return Continue +} + +func (r *box) scan( + height int, iter func(min, max []float64, value unsafe.Pointer) bool, +) bool { + n := (*node)(r.data) + if height == 0 { + for i := 0; i < n.count; i++ { + if !iter(n.boxes[i].min[:], n.boxes[i].max[:], n.boxes[i].data) { + return false + } + } + } else { + for i := 0; i < n.count; i++ { + if !n.boxes[i].scan(height-1, iter) { + return false + } + } + } + return true +} + +// Scan iterates through all items in tree. +func (tr *BoxTree) Scan(iter func(min, max []float64, value unsafe.Pointer) bool) { + if tr.root.data == nil { + return + } + tr.root.scan(tr.height, iter) +} + +// Delete ... +func (tr *BoxTree) Delete(min, max []float64, value unsafe.Pointer) { + var item box + fit(min, max, value, &item) + if tr.root.data == nil || !tr.root.contains(&item) { + return + } + var removed, recalced bool + removed, recalced, tr.reinsert = + tr.root.delete(&item, tr.height, tr.reinsert[:0]) + if !removed { + return + } + tr.count -= len(tr.reinsert) + 1 + if tr.count == 0 { + tr.root = box{} + recalced = false + } else { + for tr.height > 0 && (*node)(tr.root.data).count == 1 { + tr.root = (*node)(tr.root.data).boxes[0] + tr.height-- + tr.root.recalc() + } + } + if recalced { + tr.root.recalc() + } + for i := range tr.reinsert { + tr.insert(&tr.reinsert[i]) + tr.reinsert[i].data = nil + } +} + +func (r *box) delete(item *box, height int, reinsert []box) ( + removed, recalced bool, reinsertOut []box, +) { + n := (*node)(r.data) + if height == 0 { + for i := 0; i < n.count; i++ { + if n.boxes[i].data == item.data { + // found the target item to delete + recalced = r.onEdge(&n.boxes[i]) + n.boxes[i] = n.boxes[n.count-1] + n.boxes[n.count-1].data = nil + n.count-- + if recalced { + r.recalc() + } + return true, recalced, reinsert + } + } + } else { + for i := 0; i < n.count; i++ { + if !n.boxes[i].contains(item) { + continue + } + removed, recalced, reinsert = + n.boxes[i].delete(item, height-1, reinsert) + if !removed { + continue + } + if (*node)(n.boxes[i].data).count < minEntries { + // underflow + if !recalced { + recalced = r.onEdge(&n.boxes[i]) + } + reinsert = n.boxes[i].flatten(reinsert, height-1) + n.boxes[i] = n.boxes[n.count-1] + n.boxes[n.count-1].data = nil + n.count-- + } + if recalced { + r.recalc() + } + return removed, recalced, reinsert + } + } + return false, false, reinsert +} + +// flatten flattens all leaf boxes into a single list +func (r *box) flatten(all []box, height int) []box { + n := (*node)(r.data) + if height == 0 { + all = append(all, n.boxes[:n.count]...) + } else { + for i := 0; i < n.count; i++ { + all = n.boxes[i].flatten(all, height-1) + } + } + return all +} + +// onedge returns true when b is on the edge of r +func (r *box) onEdge(b *box) bool { + for i := 0; i < dims; i++ { + if r.min[i] == b.min[i] || r.max[i] == b.max[i] { + return true + } + } + return false +} + +// Count ... +func (tr *BoxTree) Count() int { + return tr.count +} + +func (r *box) totalOverlapArea(height int) float64 { + var area float64 + n := (*node)(r.data) + for i := 0; i < n.count; i++ { + for j := i + 1; j < n.count; j++ { + area += n.boxes[i].overlapArea(&n.boxes[j]) + } + + } + if height > 0 { + for i := 0; i < n.count; i++ { + area += n.boxes[i].totalOverlapArea(height - 1) + } + } + return area +} + +// TotalOverlapArea ... +func (tr *BoxTree) TotalOverlapArea() float64 { + if tr.root.data == nil { + return 0 + } + return tr.root.totalOverlapArea(tr.height) +} + +type qnode struct { + dist float64 + box box + height int +} + +type queue struct { + nodes []qnode + len int + size int +} + +func (q *queue) push(dist float64, box box, height int) { + if q.nodes == nil { + q.nodes = make([]qnode, 2) + } else { + q.nodes = append(q.nodes, qnode{}) + } + i := q.len + 1 + j := i / 2 + for i > 1 && q.nodes[j].dist > dist { + q.nodes[i] = q.nodes[j] + i = j + j = j / 2 + } + q.nodes[i].dist = dist + q.nodes[i].box = box + q.nodes[i].height = height + q.len++ +} + +func (q *queue) peek() qnode { + if q.len == 0 { + return qnode{} + } + return q.nodes[1] +} + +func (q *queue) pop() qnode { + if q.len == 0 { + return qnode{} + } + n := q.nodes[1] + q.nodes[1] = q.nodes[q.len] + q.len-- + var j, k int + i := 1 + for i != q.len+1 { + k = q.len + 1 + j = 2 * i + if j <= q.len && q.nodes[j].dist < q.nodes[k].dist { + k = j + } + if j+1 <= q.len && q.nodes[j+1].dist < q.nodes[k].dist { + k = j + 1 + } + q.nodes[i] = q.nodes[k] + i = k + } + return n +} + +// Nearby returns items nearest to farthest. +// The dist param is the "box distance". +func (tr *BoxTree) Nearby(min, max []float64, + iter func(min, max []float64, item unsafe.Pointer) bool) { + if tr.root.data == nil { + return + } + height := tr.height + var bbox box + fit(min, max, nil, &bbox) + var q queue + box := tr.root + for { + n := (*node)(box.data) + for i := 0; i < n.count; i++ { + dist := boxDist(&bbox, &n.boxes[i]) + q.push(dist, n.boxes[i], height-1) + } + for q.len > 0 { + if q.peek().height > -1 { + break + } + item := q.pop() + if !iter(item.box.min[:], item.box.max[:], item.box.data) { + return + } + } + if q.len == 0 { + break + } else { + qitem := q.pop() + box = qitem.box + height = qitem.height + } + } +} + +func boxDist(a, b *box) float64 { + var dist float64 + for i := 0; i < len(a.min); i++ { + var min, max float64 + if a.min[i] > b.min[i] { + min = a.min[i] + } else { + min = b.min[i] + } + if a.max[i] < b.max[i] { + max = a.max[i] + } else { + max = b.max[i] + } + squared := min - max + if squared > 0 { + dist += squared * squared + } + } + return dist +} + +// Bounds returns the minimum bounding box +func (tr *BoxTree) Bounds() (min, max []float64) { + if tr.root.data == nil { + return + } + return tr.root.min[:], tr.root.max[:] +} diff --git a/internal/collection/ptrrtree/rtree_test.go b/internal/collection/ptrrtree/rtree_test.go new file mode 100644 index 00000000..c7db757e --- /dev/null +++ b/internal/collection/ptrrtree/rtree_test.go @@ -0,0 +1,383 @@ +package ptrrtree + +import ( + "fmt" + "math/rand" + "sort" + "strconv" + "strings" + "testing" + "time" + "unsafe" +) + +type tBox struct { + min [dims]float64 + max [dims]float64 +} + +var boxes []*tBox +var points []*tBox + +func init() { + seed := time.Now().UnixNano() + // seed = 1532132365683340889 + println("seed:", seed) + rand.Seed(seed) +} + +func randPoints(N int) []*tBox { + boxes := make([]*tBox, N) + for i := 0; i < N; i++ { + boxes[i] = new(tBox) + boxes[i].min[0] = rand.Float64()*360 - 180 + boxes[i].min[1] = rand.Float64()*180 - 90 + for j := 2; j < dims; j++ { + boxes[i].min[j] = rand.Float64() + } + boxes[i].max = boxes[i].min + } + return boxes +} + +func randBoxes(N int) []*tBox { + boxes := make([]*tBox, N) + for i := 0; i < N; i++ { + boxes[i] = new(tBox) + boxes[i].min[0] = rand.Float64()*360 - 180 + boxes[i].min[1] = rand.Float64()*180 - 90 + for j := 2; j < dims; j++ { + boxes[i].min[j] = rand.Float64() * 100 + } + boxes[i].max[0] = boxes[i].min[0] + rand.Float64() + boxes[i].max[1] = boxes[i].min[1] + rand.Float64() + for j := 2; j < dims; j++ { + boxes[i].max[j] = boxes[i].min[j] + rand.Float64() + } + if boxes[i].max[0] > 180 || boxes[i].max[1] > 90 { + i-- + } + } + return boxes +} + +func sortBoxes(boxes []*tBox) { + sort.Slice(boxes, func(i, j int) bool { + for k := 0; k < len(boxes[i].min); k++ { + if boxes[i].min[k] < boxes[j].min[k] { + return true + } + if boxes[i].min[k] > boxes[j].min[k] { + return false + } + if boxes[i].max[k] < boxes[j].max[k] { + return true + } + if boxes[i].max[k] > boxes[j].max[k] { + return false + } + } + return i < j + }) +} + +func sortBoxesNearby(boxes []tBox, min, max []float64) { + sort.Slice(boxes, func(i, j int) bool { + return testBoxDist(boxes[i].min[:], boxes[i].max[:], min, max) < + testBoxDist(boxes[j].min[:], boxes[j].max[:], min, max) + }) +} + +func testBoxDist(amin, amax, bmin, bmax []float64) float64 { + var dist float64 + for i := 0; i < len(amin); i++ { + var min, max float64 + if amin[i] > bmin[i] { + min = amin[i] + } else { + min = bmin[i] + } + if amax[i] < bmax[i] { + max = amax[i] + } else { + max = bmax[i] + } + squared := min - max + if squared > 0 { + dist += squared * squared + } + } + return dist +} + +func testBoxesVarious(t *testing.T, boxes []*tBox, label string) { + N := len(boxes) + + var tr BoxTree + + // N := 10000 + // boxes := randPoints(N) + + ///////////////////////////////////////// + // insert + ///////////////////////////////////////// + for i := 0; i < N; i++ { + tr.Insert(boxes[i].min[:], boxes[i].max[:], unsafe.Pointer(boxes[i])) + } + if tr.Count() != N { + t.Fatalf("expected %d, got %d", N, tr.Count()) + } + // area := tr.TotalOverlapArea() + // fmt.Printf("overlap: %.0f, %.1f/item\n", area, area/float64(N)) + + // ioutil.WriteFile(label+".svg", []byte(rtreetools.SVG(&tr)), 0600) + + ///////////////////////////////////////// + // scan all items and count one-by-one + ///////////////////////////////////////// + var count int + tr.Scan(func(min, max []float64, value unsafe.Pointer) bool { + count++ + return true + }) + if count != N { + t.Fatalf("expected %d, got %d", N, count) + } + + ///////////////////////////////////////// + // check every point for correctness + ///////////////////////////////////////// + var tboxes1 []*tBox + tr.Scan(func(min, max []float64, value unsafe.Pointer) bool { + tboxes1 = append(tboxes1, (*tBox)(value)) + return true + }) + tboxes2 := make([]*tBox, len(boxes)) + copy(tboxes2, boxes) + sortBoxes(tboxes1) + sortBoxes(tboxes2) + for i := 0; i < len(tboxes1); i++ { + if tboxes1[i] != tboxes2[i] { + t.Fatalf("expected '%v', got '%v'", tboxes2[i], tboxes1[i]) + } + } + + ///////////////////////////////////////// + // search for each item one-by-one + ///////////////////////////////////////// + for i := 0; i < N; i++ { + var found bool + tr.Search(boxes[i].min[:], boxes[i].max[:], + func(min, max []float64, value unsafe.Pointer) bool { + if value == unsafe.Pointer(boxes[i]) { + found = true + return false + } + return true + }) + if !found { + t.Fatalf("did not find item %d", i) + } + } + + centerMin, centerMax := []float64{-18, -9}, []float64{18, 9} + for j := 2; j < dims; j++ { + centerMin = append(centerMin, -10) + centerMax = append(centerMax, 10) + } + + ///////////////////////////////////////// + // search for 10% of the items + ///////////////////////////////////////// + for i := 0; i < N/5; i++ { + var count int + tr.Search(centerMin, centerMax, + func(min, max []float64, value unsafe.Pointer) bool { + count++ + return true + }, + ) + } + + ///////////////////////////////////////// + // delete every other item + ///////////////////////////////////////// + for i := 0; i < N/2; i++ { + j := i * 2 + tr.Delete(boxes[j].min[:], boxes[j].max[:], unsafe.Pointer(boxes[j])) + } + + ///////////////////////////////////////// + // count all items. should be half of N + ///////////////////////////////////////// + count = 0 + tr.Scan(func(min, max []float64, value unsafe.Pointer) bool { + count++ + return true + }) + if count != N/2 { + t.Fatalf("expected %d, got %d", N/2, count) + } + + /////////////////////////////////////////////////// + // reinsert every other item, but in random order + /////////////////////////////////////////////////// + var ij []int + for i := 0; i < N/2; i++ { + j := i * 2 + ij = append(ij, j) + } + rand.Shuffle(len(ij), func(i, j int) { + ij[i], ij[j] = ij[j], ij[i] + }) + for i := 0; i < N/2; i++ { + j := ij[i] + tr.Insert(boxes[j].min[:], boxes[j].max[:], unsafe.Pointer(boxes[j])) + } + + ////////////////////////////////////////////////////// + // replace each item with an item that is very close + ////////////////////////////////////////////////////// + var nboxes = make([]*tBox, N) + for i := 0; i < N; i++ { + nboxes[i] = new(tBox) + for j := 0; j < len(boxes[i].min); j++ { + nboxes[i].min[j] = boxes[i].min[j] + (rand.Float64() - 0.5) + if boxes[i].min == boxes[i].max { + nboxes[i].max[j] = nboxes[i].min[j] + } else { + nboxes[i].max[j] = boxes[i].max[j] + (rand.Float64() - 0.5) + } + } + + } + for i := 0; i < N; i++ { + tr.Insert(nboxes[i].min[:], nboxes[i].max[:], unsafe.Pointer(nboxes[i])) + tr.Delete(boxes[i].min[:], boxes[i].max[:], unsafe.Pointer(boxes[i])) + } + if tr.Count() != N { + t.Fatalf("expected %d, got %d", N, tr.Count()) + } + // area = tr.TotalOverlapArea() + // fmt.Fprintf(wr, "overlap: %.0f, %.1f/item\n", area, area/float64(N)) + + ///////////////////////////////////////// + // check every point for correctness + ///////////////////////////////////////// + tboxes1 = nil + tr.Scan(func(min, max []float64, value unsafe.Pointer) bool { + tboxes1 = append(tboxes1, (*tBox)(value)) + return true + }) + tboxes2 = make([]*tBox, len(nboxes)) + copy(tboxes2, nboxes) + sortBoxes(tboxes1) + sortBoxes(tboxes2) + for i := 0; i < len(tboxes1); i++ { + if tboxes1[i] != tboxes2[i] { + t.Fatalf("expected '%v', got '%v'", tboxes2[i], tboxes1[i]) + } + } + + ///////////////////////////////////////// + // search for 10% of the items + ///////////////////////////////////////// + for i := 0; i < N/5; i++ { + var count int + tr.Search(centerMin, centerMax, + func(min, max []float64, value unsafe.Pointer) bool { + count++ + return true + }, + ) + } + + var boxes3 []*tBox + tr.Nearby(centerMin, centerMax, + func(min, max []float64, value unsafe.Pointer) bool { + boxes3 = append(boxes3, (*tBox)(value)) + return true + }, + ) + if len(boxes3) != len(nboxes) { + t.Fatalf("expected %d, got %d", len(nboxes), len(boxes3)) + } + if len(boxes3) != tr.Count() { + t.Fatalf("expected %d, got %d", tr.Count(), len(boxes3)) + } + var ldist float64 + for i, box := range boxes3 { + dist := testBoxDist(box.min[:], box.max[:], centerMin, centerMax) + if i > 0 && dist < ldist { + t.Fatalf("out of order") + } + ldist = dist + } +} + +func TestRandomBoxes(t *testing.T) { + testBoxesVarious(t, randBoxes(10000), "boxes") +} + +func TestRandomPoints(t *testing.T) { + testBoxesVarious(t, randPoints(10000), "points") +} + +func (r *box) boxstr() string { + var b []byte + b = append(b, '[', '[') + for i := 0; i < len(r.min); i++ { + if i != 0 { + b = append(b, ' ') + } + b = strconv.AppendFloat(b, r.min[i], 'f', -1, 64) + } + b = append(b, ']', '[') + for i := 0; i < len(r.max); i++ { + if i != 0 { + b = append(b, ' ') + } + b = strconv.AppendFloat(b, r.max[i], 'f', -1, 64) + } + b = append(b, ']', ']') + return string(b) +} + +func (r *box) print(height, indent int) { + fmt.Printf("%s%s", strings.Repeat(" ", indent), r.boxstr()) + if height == 0 { + fmt.Printf("\t'%v'\n", r.data) + } else { + fmt.Printf("\n") + for i := 0; i < (*node)(r.data).count; i++ { + (*node)(r.data).boxes[i].print(height-1, indent+1) + } + } + +} + +func (tr BoxTree) print() { + if tr.root.data == nil { + println("EMPTY TREE") + return + } + tr.root.print(tr.height+1, 0) +} + +func TestZeroPoints(t *testing.T) { + N := 10000 + var tr BoxTree + pt := make([]float64, dims) + for i := 0; i < N; i++ { + tr.Insert(pt, nil, nil) + } +} + +func BenchmarkRandomInsert(b *testing.B) { + var tr BoxTree + boxes := randBoxes(b.N) + b.ResetTimer() + for i := 0; i < b.N; i++ { + tr.Insert(boxes[i].min[:], boxes[i].max[:], nil) + } +}