tile38/internal/collection/ptrrtree/rtree.go

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[:]
}