3
0
mirror of https://github.com/ergochat/ergo.git synced 2024-12-05 10:29:25 +01:00
ergo/vendor/github.com/tidwall/rtree/base/rtree.go

675 lines
15 KiB
Go
Raw Normal View History

package base
import (
"math"
"unsafe"
)
// precalculate infinity
var mathInfNeg = math.Inf(-1)
var mathInfPos = math.Inf(+1)
type treeNode struct {
min, max []float64
children []*treeNode
count int
height int
leaf bool
}
func (node *treeNode) unsafeItem() *treeItem {
return (*treeItem)(unsafe.Pointer(node))
}
func (tr *RTree) createNode(children []*treeNode) *treeNode {
n := &treeNode{
height: 1,
leaf: true,
children: make([]*treeNode, tr.maxEntries+1),
}
if len(children) > 0 {
n.count = len(children)
copy(n.children[:n.count], children)
}
n.min = make([]float64, tr.dims)
n.max = make([]float64, tr.dims)
for i := 0; i < tr.dims; i++ {
n.min[i] = mathInfPos
n.max[i] = mathInfNeg
}
return n
}
func (node *treeNode) extend(b *treeNode) {
for i := 0; i < len(node.min); i++ {
if b.min[i] < node.min[i] {
node.min[i] = b.min[i]
}
if b.max[i] > node.max[i] {
node.max[i] = b.max[i]
}
}
}
func (node *treeNode) area() float64 {
area := node.max[0] - node.min[0]
for i := 1; i < len(node.min); i++ {
area *= node.max[i] - node.min[i]
}
return area
}
func (node *treeNode) enlargedAreaAxis(b *treeNode, axis int) float64 {
var max, min float64
if b.max[axis] > node.max[axis] {
max = b.max[axis]
} else {
max = node.max[axis]
}
if b.min[axis] < node.min[axis] {
min = b.min[axis]
} else {
min = node.min[axis]
}
return max - min
}
func (node *treeNode) enlargedArea(b *treeNode) float64 {
area := node.enlargedAreaAxis(b, 0)
for i := 1; i < len(node.min); i++ {
area *= node.enlargedAreaAxis(b, i)
}
return area
}
func (node *treeNode) intersectionAreaAxis(b *treeNode, axis int) float64 {
var max, min float64
if node.max[axis] < b.max[axis] {
max = node.max[axis]
} else {
max = b.max[axis]
}
if node.min[axis] > b.min[axis] {
min = node.min[axis]
} else {
min = b.min[axis]
}
if max > min {
return max - min
}
return 0
}
func (node *treeNode) intersectionArea(b *treeNode) float64 {
area := node.intersectionAreaAxis(b, 0)
for i := 1; i < len(node.min); i++ {
area *= node.intersectionAreaAxis(b, i)
}
return area
}
func (node *treeNode) margin() float64 {
margin := node.max[0] - node.min[0]
for i := 1; i < len(node.min); i++ {
margin += node.max[i] - node.min[i]
}
return margin
}
type result int
const (
not result = 0
intersects result = 1
contains result = 2
)
func (node *treeNode) overlaps(b *treeNode) result {
for i := 0; i < len(node.min); i++ {
if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
return not
}
if node.min[i] > b.min[i] || b.max[i] > node.max[i] {
i++
for ; i < len(node.min); i++ {
if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
return not
}
}
return intersects
}
}
return contains
}
func (node *treeNode) intersects(b *treeNode) bool {
for i := 0; i < len(node.min); i++ {
if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
return false
}
}
return true
}
func (node *treeNode) findItem(item interface{}) int {
for i := 0; i < node.count; i++ {
if node.children[i].unsafeItem().item == item {
return i
}
}
return -1
}
func (node *treeNode) contains(b *treeNode) bool {
for i := 0; i < len(node.min); i++ {
if node.min[i] > b.min[i] || b.max[i] > node.max[i] {
return false
}
}
return true
}
func (node *treeNode) childCount() int {
if node.leaf {
return node.count
}
var n int
for i := 0; i < node.count; i++ {
n += node.children[i].childCount()
}
return n
}
type treeItem struct {
min, max []float64
item interface{}
}
2020-11-08 23:55:22 +01:00
//go:nocheckptr
func (item *treeItem) unsafeNode() *treeNode {
return (*treeNode)(unsafe.Pointer(item))
}
// RTree is an R-tree
type RTree struct {
dims int
maxEntries int
minEntries int
data *treeNode // root node
// resusable fields, these help performance of common mutable operations.
reuse struct {
path []*treeNode // for reinsertion path
indexes []int // for remove function
stack []int // for bulk loading
}
}
// New creates a new R-tree
func New(dims, maxEntries int) *RTree {
if dims <= 0 {
panic("invalid dimensions")
}
tr := &RTree{}
tr.dims = dims
tr.maxEntries = int(math.Max(4, float64(maxEntries)))
tr.minEntries = int(math.Max(2, math.Ceil(float64(tr.maxEntries)*0.4)))
tr.data = tr.createNode(nil)
return tr
}
// Insert inserts an item
func (tr *RTree) Insert(min, max []float64, item interface{}) {
if len(min) != tr.dims || len(max) != tr.dims {
panic("invalid dimensions")
}
if item == nil {
panic("nil item")
}
bbox := treeNode{min: min, max: max}
tr.insert(&bbox, item, tr.data.height-1, false)
}
func (tr *RTree) insert(bbox *treeNode, item interface{}, level int, isNode bool) {
tr.reuse.path = tr.reuse.path[:0]
node, insertPath := tr.chooseSubtree(bbox, tr.data, level, tr.reuse.path)
if item == nil {
// item is only nil when bulk loading a node
if node.leaf {
panic("loading node into leaf")
}
node.children[node.count] = bbox
node.count++
} else {
ti := &treeItem{min: bbox.min, max: bbox.max, item: item}
node.children[node.count] = ti.unsafeNode()
node.count++
}
node.extend(bbox)
for level >= 0 {
if insertPath[level].count > tr.maxEntries {
insertPath = tr.split(insertPath, level)
level--
} else {
break
}
}
tr.adjustParentBBoxes(bbox, insertPath, level)
tr.reuse.path = insertPath
}
func (tr *RTree) adjustParentBBoxes(bbox *treeNode, path []*treeNode, level int) {
// adjust bboxes along the given tree path
for i := level; i >= 0; i-- {
path[i].extend(bbox)
}
}
func (tr *RTree) chooseSubtree(bbox, node *treeNode, level int, path []*treeNode) (*treeNode, []*treeNode) {
var targetNode *treeNode
var area, enlargement, minArea, minEnlargement float64
for {
path = append(path, node)
if node.leaf || len(path)-1 == level {
break
}
minEnlargement = mathInfPos
minArea = minEnlargement
for i := 0; i < node.count; i++ {
child := node.children[i]
area = child.area()
enlargement = bbox.enlargedArea(child) - area
if enlargement < minEnlargement {
minEnlargement = enlargement
if area < minArea {
minArea = area
}
targetNode = child
} else if enlargement == minEnlargement {
if area < minArea {
minArea = area
targetNode = child
}
}
}
if targetNode != nil {
node = targetNode
} else if node.count > 0 {
node = (*treeNode)(node.children[0])
} else {
node = nil
}
}
return node, path
}
func (tr *RTree) split(insertPath []*treeNode, level int) []*treeNode {
var node = insertPath[level]
var M = node.count
var m = tr.minEntries
tr.chooseSplitAxis(node, m, M)
splitIndex := tr.chooseSplitIndex(node, m, M)
spliced := make([]*treeNode, node.count-splitIndex)
copy(spliced, node.children[splitIndex:])
node.count = splitIndex
newNode := tr.createNode(spliced)
newNode.height = node.height
newNode.leaf = node.leaf
tr.calcBBox(node)
tr.calcBBox(newNode)
if level != 0 {
insertPath[level-1].children[insertPath[level-1].count] = newNode
insertPath[level-1].count++
} else {
tr.splitRoot(node, newNode)
}
return insertPath
}
func (tr *RTree) chooseSplitIndex(node *treeNode, m, M int) int {
var i int
var bbox1, bbox2 *treeNode
var overlap, area, minOverlap, minArea float64
var index int
minArea = mathInfPos
minOverlap = minArea
for i = m; i <= M-m; i++ {
bbox1 = tr.distBBox(node, 0, i, nil)
bbox2 = tr.distBBox(node, i, M, nil)
overlap = bbox1.intersectionArea(bbox2)
area = bbox1.area() + bbox2.area()
// choose distribution with minimum overlap
if overlap < minOverlap {
minOverlap = overlap
index = i
if area < minArea {
minArea = area
}
} else if overlap == minOverlap {
// otherwise choose distribution with minimum area
if area < minArea {
minArea = area
index = i
}
}
}
return index
}
func (tr *RTree) calcBBox(node *treeNode) {
tr.distBBox(node, 0, node.count, node)
}
func (tr *RTree) chooseSplitAxis(node *treeNode, m, M int) {
minMargin := tr.allDistMargin(node, m, M, 0)
var minAxis int
for axis := 1; axis < tr.dims; axis++ {
margin := tr.allDistMargin(node, m, M, axis)
if margin < minMargin {
minMargin = margin
minAxis = axis
}
}
if minAxis < tr.dims {
tr.sortNodes(node, minAxis)
}
}
func (tr *RTree) splitRoot(node, newNode *treeNode) {
tr.data = tr.createNode([]*treeNode{node, newNode})
tr.data.height = node.height + 1
tr.data.leaf = false
tr.calcBBox(tr.data)
}
func (tr *RTree) distBBox(node *treeNode, k, p int, destNode *treeNode) *treeNode {
if destNode == nil {
destNode = tr.createNode(nil)
} else {
for i := 0; i < tr.dims; i++ {
destNode.min[i] = mathInfPos
destNode.max[i] = mathInfNeg
}
}
for i := k; i < p; i++ {
if node.leaf {
destNode.extend(node.children[i])
} else {
destNode.extend((*treeNode)(node.children[i]))
}
}
return destNode
}
func (tr *RTree) allDistMargin(node *treeNode, m, M int, axis int) float64 {
tr.sortNodes(node, axis)
var leftBBox = tr.distBBox(node, 0, m, nil)
var rightBBox = tr.distBBox(node, M-m, M, nil)
var margin = leftBBox.margin() + rightBBox.margin()
var i int
if node.leaf {
for i = m; i < M-m; i++ {
leftBBox.extend(node.children[i])
margin += leftBBox.margin()
}
for i = M - m - 1; i >= m; i-- {
leftBBox.extend(node.children[i])
margin += rightBBox.margin()
}
} else {
for i = m; i < M-m; i++ {
child := (*treeNode)(node.children[i])
leftBBox.extend(child)
margin += leftBBox.margin()
}
for i = M - m - 1; i >= m; i-- {
child := (*treeNode)(node.children[i])
leftBBox.extend(child)
margin += rightBBox.margin()
}
}
return margin
}
func (tr *RTree) sortNodes(node *treeNode, axis int) {
sortByAxis(node.children[:node.count], axis)
}
func sortByAxis(items []*treeNode, axis int) {
if len(items) < 2 {
return
}
left, right := 0, len(items)-1
pivotIndex := len(items) / 2
items[pivotIndex], items[right] = items[right], items[pivotIndex]
for i := range items {
if items[i].min[axis] < items[right].min[axis] {
items[i], items[left] = items[left], items[i]
left++
}
}
items[left], items[right] = items[right], items[left]
sortByAxis(items[:left], axis)
sortByAxis(items[left+1:], axis)
}
// Search searches the tree for items in the input rectangle
func (tr *RTree) Search(min, max []float64, iter func(item interface{}) bool) bool {
bbox := &treeNode{min: min, max: max}
if !tr.data.intersects(bbox) {
return true
}
return tr.search(tr.data, bbox, iter)
}
func (tr *RTree) search(node, bbox *treeNode, iter func(item interface{}) bool) bool {
if node.leaf {
for i := 0; i < node.count; i++ {
if bbox.intersects(node.children[i]) {
if !iter(node.children[i].unsafeItem().item) {
return false
}
}
}
} else {
for i := 0; i < node.count; i++ {
r := bbox.overlaps(node.children[i])
if r == intersects {
if !tr.search(node.children[i], bbox, iter) {
return false
}
} else if r == contains {
if !scan(node.children[i], iter) {
return false
}
}
}
}
return true
}
func (tr *RTree) IsEmpty() bool {
empty := true
tr.Scan(func(item interface{}) bool {
empty = false
return false
})
return empty
}
// Remove removes an item from the R-tree.
func (tr *RTree) Remove(min, max []float64, item interface{}) {
bbox := &treeNode{min: min, max: max}
tr.remove(bbox, item)
}
func (tr *RTree) remove(bbox *treeNode, item interface{}) {
path := tr.reuse.path[:0]
indexes := tr.reuse.indexes[:0]
var node = tr.data
var i int
var parent *treeNode
var index int
var goingUp bool
for node != nil || len(path) != 0 {
if node == nil {
node = path[len(path)-1]
path = path[:len(path)-1]
if len(path) == 0 {
parent = nil
} else {
parent = path[len(path)-1]
}
i = indexes[len(indexes)-1]
indexes = indexes[:len(indexes)-1]
goingUp = true
}
if node.leaf {
index = node.findItem(item)
if index != -1 {
// item found, remove the item and condense tree upwards
copy(node.children[index:], node.children[index+1:])
node.children[node.count-1] = nil
node.count--
path = append(path, node)
tr.condense(path)
goto done
}
}
if !goingUp && !node.leaf && node.contains(bbox) { // go down
path = append(path, node)
indexes = append(indexes, i)
i = 0
parent = node
node = (*treeNode)(node.children[0])
} else if parent != nil { // go right
i++
if i == parent.count {
node = nil
} else {
node = (*treeNode)(parent.children[i])
}
goingUp = false
} else {
node = nil
}
}
done:
tr.reuse.path = path
tr.reuse.indexes = indexes
return
}
func (tr *RTree) condense(path []*treeNode) {
// go through the path, removing empty nodes and updating bboxes
var siblings []*treeNode
for i := len(path) - 1; i >= 0; i-- {
if path[i].count == 0 {
if i > 0 {
siblings = path[i-1].children[:path[i-1].count]
index := -1
for j := 0; j < len(siblings); j++ {
if siblings[j] == path[i] {
index = j
break
}
}
copy(siblings[index:], siblings[index+1:])
siblings[len(siblings)-1] = nil
path[i-1].count--
//siblings = siblings[:len(siblings)-1]
//path[i-1].children = siblings
} else {
tr.data = tr.createNode(nil) // clear tree
}
} else {
tr.calcBBox(path[i])
}
}
}
// Count returns the number of items in the R-tree.
func (tr *RTree) Count() int {
return tr.data.childCount()
}
// Traverse iterates over the entire R-tree and includes all nodes and items.
func (tr *RTree) Traverse(iter func(min, max []float64, level int, item interface{}) bool) bool {
return tr.traverse(tr.data, iter)
}
func (tr *RTree) traverse(node *treeNode, iter func(min, max []float64, level int, item interface{}) bool) bool {
if !iter(node.min, node.max, int(node.height), nil) {
return false
}
if node.leaf {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !iter(child.min, child.max, 0, child.unsafeItem().item) {
return false
}
}
} else {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !tr.traverse(child, iter) {
return false
}
}
}
return true
}
// Scan iterates over the entire R-tree
func (tr *RTree) Scan(iter func(item interface{}) bool) bool {
return scan(tr.data, iter)
}
func scan(node *treeNode, iter func(item interface{}) bool) bool {
if node.leaf {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !iter(child.unsafeItem().item) {
return false
}
}
} else {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !scan(child, iter) {
return false
}
}
}
return true
}
// Bounds returns the bounding box of the entire R-tree
func (tr *RTree) Bounds() (min, max []float64) {
if tr.data.count > 0 {
return tr.data.min, tr.data.max
}
return make([]float64, tr.dims), make([]float64, tr.dims)
}
// Complexity returns the complexity of the R-tree. The higher the value, the
// more complex the tree. The value of 1 is the lowest.
func (tr *RTree) Complexity() float64 {
var nodeCount int
var itemCount int
tr.Traverse(func(_, _ []float64, level int, _ interface{}) bool {
if level == 0 {
itemCount++
} else {
nodeCount++
}
return true
})
return float64(tr.maxEntries*nodeCount) / float64(itemCount)
}