rtred/rtree.go

279 lines
5.3 KiB
Go
Raw Normal View History

2016-06-12 02:43:14 +03:00
package rtree
2018-01-13 02:31:15 +03:00
import (
"math"
"sync"
"github.com/tidwall/rtree/base"
)
2016-06-12 02:43:14 +03:00
type Iterator func(item Item) bool
type Item interface {
2016-07-18 03:47:23 +03:00
Rect(ctx interface{}) (min []float64, max []float64)
2016-06-12 02:43:14 +03:00
}
type RTree struct {
2018-01-13 02:31:15 +03:00
dims int
maxEntries int
ctx interface{}
trs []*base.RTree
used int
2016-06-12 02:43:14 +03:00
}
2016-07-18 03:47:23 +03:00
func New(ctx interface{}) *RTree {
2018-01-13 02:31:15 +03:00
tr := &RTree{
ctx: ctx,
dims: 20,
maxEntries: 13,
2016-06-12 02:43:14 +03:00
}
2018-01-13 02:31:15 +03:00
tr.trs = make([]*base.RTree, 20)
return tr
2016-06-12 02:43:14 +03:00
}
func (tr *RTree) Insert(item Item) {
if item == nil {
2018-01-13 02:31:15 +03:00
panic("nil item")
2016-06-12 02:43:14 +03:00
}
2016-07-18 03:47:23 +03:00
min, max := item.Rect(tr.ctx)
2016-06-12 02:43:14 +03:00
if len(min) != len(max) {
2016-07-18 15:26:26 +03:00
return // just return
2016-06-12 02:43:14 +03:00
panic("invalid item rectangle")
}
2018-01-13 02:31:15 +03:00
if len(min) < 1 || len(min) > len(tr.trs) {
2016-07-18 15:26:26 +03:00
return // just return
2016-06-12 02:43:14 +03:00
panic("invalid dimension")
}
2018-01-13 02:31:15 +03:00
btr := tr.trs[len(min)-1]
if btr == nil {
btr = base.New(len(min), tr.maxEntries)
tr.trs[len(min)-1] = btr
tr.used++
}
amin := make([]float64, len(min))
amax := make([]float64, len(max))
for i := 0; i < len(min); i++ {
amin[i], amax[i] = min[i], max[i]
}
btr.Insert(amin, amax, item)
2016-06-12 02:43:14 +03:00
}
func (tr *RTree) Remove(item Item) {
if item == nil {
2018-01-13 02:31:15 +03:00
panic("nil item")
2016-06-12 02:43:14 +03:00
}
2016-07-18 03:47:23 +03:00
min, max := item.Rect(tr.ctx)
2016-06-12 02:43:14 +03:00
if len(min) != len(max) {
2016-07-18 15:26:26 +03:00
return // just return
2016-06-12 02:43:14 +03:00
panic("invalid item rectangle")
}
2018-01-13 02:31:15 +03:00
if len(min) < 1 || len(min) > len(tr.trs) {
2016-07-18 15:26:26 +03:00
return // just return
2016-06-12 02:43:14 +03:00
panic("invalid dimension")
2018-01-13 02:31:15 +03:00
}
btr := tr.trs[len(min)-1]
if btr == nil {
return
}
amin := make([]float64, len(min))
amax := make([]float64, len(max))
for i := 0; i < len(min); i++ {
amin[i], amax[i] = min[i], max[i]
}
btr.Remove(amin, amax, item)
if btr.IsEmpty() {
tr.trs[len(min)-1] = nil
tr.used--
2016-06-12 02:43:14 +03:00
}
}
func (tr *RTree) Reset() {
2018-01-13 02:31:15 +03:00
for i := 0; i < len(tr.trs); i++ {
tr.trs[i] = nil
}
tr.used = 0
2016-06-12 02:43:14 +03:00
}
func (tr *RTree) Count() int {
2018-01-13 02:31:15 +03:00
var count int
for _, btr := range tr.trs {
if btr != nil {
count += btr.Count()
}
}
2016-07-22 01:30:24 +03:00
return count
2016-06-12 02:43:14 +03:00
}
2018-01-13 02:31:15 +03:00
2016-06-12 02:43:14 +03:00
func (tr *RTree) Search(bounds Item, iter Iterator) {
if bounds == nil {
2016-07-18 15:26:26 +03:00
panic("nil bounds being used for search")
2016-06-12 02:43:14 +03:00
}
2016-07-18 03:47:23 +03:00
min, max := bounds.Rect(tr.ctx)
2016-06-12 02:43:14 +03:00
if len(min) != len(max) {
2016-07-18 15:26:26 +03:00
return // just return
2016-06-12 02:43:14 +03:00
panic("invalid item rectangle")
}
2018-01-13 02:31:15 +03:00
if len(min) < 1 || len(min) > len(tr.trs) {
2016-07-18 15:26:26 +03:00
return // just return
2016-06-12 02:43:14 +03:00
panic("invalid dimension")
}
2018-01-13 02:31:15 +03:00
used := tr.used
for i, btr := range tr.trs {
if used == 0 {
break
2016-06-12 02:43:14 +03:00
}
2018-01-13 02:31:15 +03:00
if btr != nil {
if !search(btr, min, max, i+1, iter) {
return
}
used--
2016-06-12 02:43:14 +03:00
}
}
}
2018-01-13 02:31:15 +03:00
func search(btr *base.RTree, min, max []float64, dims int, iter Iterator) bool {
amin := make([]float64, dims)
amax := make([]float64, dims)
for i := 0; i < dims; i++ {
2016-07-22 01:30:24 +03:00
if i < len(min) {
amin[i] = min[i]
amax[i] = max[i]
} else {
amin[i] = math.Inf(-1)
amax[i] = math.Inf(+1)
}
}
2018-01-13 02:31:15 +03:00
var ended bool
btr.Search(amin, amax, func(item interface{}) bool {
if !iter(item.(Item)) {
2016-07-22 01:30:24 +03:00
ended = true
return false
}
return true
})
return !ended
}
2018-01-13 02:31:15 +03:00
func (tr *RTree) KNN(bounds Item, center bool, iter func(item Item, dist float64) bool) {
if bounds == nil {
panic("nil bounds being used for search")
2016-07-22 01:30:24 +03:00
}
2018-01-13 02:31:15 +03:00
min, max := bounds.Rect(tr.ctx)
if len(min) != len(max) {
return // just return
panic("invalid item rectangle")
2016-07-22 01:30:24 +03:00
}
2018-01-13 02:31:15 +03:00
if len(min) < 1 || len(min) > len(tr.trs) {
return // just return
panic("invalid dimension")
2016-07-22 01:30:24 +03:00
}
2018-01-13 02:31:15 +03:00
if tr.used == 0 {
return
2016-07-22 01:30:24 +03:00
}
2018-01-13 02:31:15 +03:00
if tr.used == 1 {
for i, btr := range tr.trs {
if btr != nil {
knn(btr, min, max, center, i+1, func(item interface{}, dist float64) bool {
return iter(item.(Item), dist)
})
break
}
2016-07-22 01:30:24 +03:00
}
2018-01-13 02:31:15 +03:00
return
2016-07-22 01:30:24 +03:00
}
2018-01-13 02:31:15 +03:00
type queueT struct {
done bool
step int
item Item
dist float64
}
var mu sync.Mutex
var ended bool
queues := make(map[int][]queueT)
cond := sync.NewCond(&mu)
for i, btr := range tr.trs {
if btr != nil {
dims := i + 1
mu.Lock()
queues[dims] = []queueT{}
cond.Signal()
mu.Unlock()
go func(dims int, btr *base.RTree) {
knn(btr, min, max, center, dims, func(item interface{}, dist float64) bool {
mu.Lock()
if ended {
mu.Unlock()
return false
}
queues[dims] = append(queues[dims], queueT{item: item.(Item), dist: dist})
cond.Signal()
mu.Unlock()
return true
})
mu.Lock()
queues[dims] = append(queues[dims], queueT{done: true})
cond.Signal()
mu.Unlock()
}(dims, btr)
}
}
mu.Lock()
for {
ready := true
for i := range queues {
if len(queues[i]) == 0 {
ready = false
break
}
if queues[i][0].done {
delete(queues, i)
}
}
if len(queues) == 0 {
break
}
if ready {
var j int
var minDist float64
var minItem Item
var minQueue int
for i := range queues {
if j == 0 || queues[i][0].dist < minDist {
minDist = queues[i][0].dist
minItem = queues[i][0].item
minQueue = i
}
}
queues[minQueue] = queues[minQueue][1:]
if !iter(minItem, minDist) {
ended = true
break
}
continue
}
cond.Wait()
}
mu.Unlock()
}
func knn(btr *base.RTree, min, max []float64, center bool, dims int, iter func(item interface{}, dist float64) bool) bool {
amin := make([]float64, dims)
amax := make([]float64, dims)
for i := 0; i < dims; i++ {
2016-07-22 01:30:24 +03:00
if i < len(min) {
amin[i] = min[i]
amax[i] = max[i]
} else {
amin[i] = math.Inf(-1)
amax[i] = math.Inf(+1)
}
}
2018-01-13 02:31:15 +03:00
var ended bool
btr.KNN(amin, amax, center, func(item interface{}, dist float64) bool {
if !iter(item.(Item), dist) {
2016-07-22 01:30:24 +03:00
ended = true
return false
}
return true
})
return !ended
}