tile38/vendor/github.com/tidwall/geojson/geometry/rtree.go

package geometry

import (
	"encoding/binary"
	"math"
)

const rDims = 2
const rMaxEntries = 16

type rRect struct {
	data     interface{}
	min, max [rDims]float64
}

type rNode struct {
	count int
	rects [rMaxEntries + 1]rRect
}

// rTree ...
type rTree struct {
	height   int
	root     rRect
	count    int
	reinsert []rRect
}

func (r *rRect) expand(b *rRect) {
	for i := 0; i < rDims; 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]
		}
	}
}

// Insert inserts an item into the RTree
func (tr *rTree) Insert(min, max []float64, value interface{}) {
	var item rRect
	fit(min, max, value, &item)
	tr.insert(&item)
}

func (tr *rTree) insert(item *rRect) {
	if tr.root.data == nil {
		fit(item.min[:], item.max[:], new(rNode), &tr.root)
	}
	grown := tr.root.insert(item, tr.height)
	if grown {
		tr.root.expand(item)
	}
	if tr.root.data.(*rNode).count == rMaxEntries+1 {
		newRoot := new(rNode)
		tr.root.splitLargestAxisEdgeSnap(&newRoot.rects[1])
		newRoot.rects[0] = tr.root
		newRoot.count = 2
		tr.root.data = newRoot
		tr.root.recalc()
		tr.height++
	}
	tr.count++
}

func (r *rRect) chooseLeastEnlargement(b *rRect) int {
	j, jenlargement, jarea := -1, 0.0, 0.0
	n := r.data.(*rNode)
	for i := 0; i < n.count; i++ {
		// force inline
		area := n.rects[i].max[0] - n.rects[i].min[0]
		for j := 1; j < rDims; j++ {
			area *= n.rects[i].max[j] - n.rects[i].min[j]
		}
		var enlargement float64
		// force inline
		enlargedArea := 1.0
		for j := 0; j < len(n.rects[i].min); j++ {
			if b.max[j] > n.rects[i].max[j] {
				if b.min[j] < n.rects[i].min[j] {
					enlargedArea *= b.max[j] - b.min[j]
				} else {
					enlargedArea *= b.max[j] - n.rects[i].min[j]
				}
			} else {
				if b.min[j] < n.rects[i].min[j] {
					enlargedArea *= n.rects[i].max[j] - b.min[j]
				} else {
					enlargedArea *= n.rects[i].max[j] - n.rects[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 *rRect) recalc() {
	n := r.data.(*rNode)
	r.min = n.rects[0].min
	r.max = n.rects[0].max
	for i := 1; i < n.count; i++ {
		r.expand(&n.rects[i])
	}
}

// contains return struct when b is fully contained inside of n
func (r *rRect) contains(b *rRect) bool {
	for i := 0; i < rDims; i++ {
		if b.min[i] < r.min[i] || b.max[i] > r.max[i] {
			return false
		}
	}
	return true
}

func (r *rRect) largestAxis() (axis int, size float64) {
	j, jsz := 0, 0.0
	for i := 0; i < rDims; i++ {
		sz := r.max[i] - r.min[i]
		if i == 0 || sz > jsz {
			j, jsz = i, sz
		}
	}
	return j, jsz
}

func (r *rRect) splitLargestAxisEdgeSnap(right *rRect) {
	axis, _ := r.largestAxis()
	left := r
	leftNode := left.data.(*rNode)
	rightNode := new(rNode)
	right.data = rightNode

	var equals []rRect
	for i := 0; i < leftNode.count; i++ {
		minDist := leftNode.rects[i].min[axis] - left.min[axis]
		maxDist := left.max[axis] - leftNode.rects[i].max[axis]
		if minDist < maxDist {
			// stay left
		} else {
			if minDist > maxDist {
				// move to right
				rightNode.rects[rightNode.count] = leftNode.rects[i]
				rightNode.count++
			} else {
				// move to equals, at the end of the left array
				equals = append(equals, leftNode.rects[i])
			}
			leftNode.rects[i] = leftNode.rects[leftNode.count-1]
			leftNode.rects[leftNode.count-1].data = nil
			leftNode.count--
			i--
		}
	}
	for _, b := range equals {
		if leftNode.count < rightNode.count {
			leftNode.rects[leftNode.count] = b
			leftNode.count++
		} else {
			rightNode.rects[rightNode.count] = b
			rightNode.count++
		}
	}
	left.recalc()
	right.recalc()
}

func (r *rRect) insert(item *rRect, height int) (grown bool) {
	n := r.data.(*rNode)
	if height == 0 {
		n.rects[n.count] = *item
		n.count++
		grown = !r.contains(item)
		return grown
	}
	// choose subtree
	index := r.chooseLeastEnlargement(item)
	child := &n.rects[index]
	grown = child.insert(item, height-1)
	if grown {
		child.expand(item)
		grown = !r.contains(item)
	}
	if child.data.(*rNode).count == rMaxEntries+1 {
		child.splitLargestAxisEdgeSnap(&n.rects[n.count])
		n.count++
	}
	return grown
}

// fit an external item into a rect type
func fit(min, max []float64, value interface{}, target *rRect) {
	if max == nil {
		max = min
	}
	if len(min) != len(max) {
		panic("min/max dimension mismatch")
	}
	if len(min) != rDims {
		panic("invalid number of dimensions")
	}
	for i := 0; i < rDims; i++ {
		target.min[i] = min[i]
		target.max[i] = max[i]
	}
	target.data = value
}

func (r *rRect) intersects(b *rRect) bool {
	for i := 0; i < rDims; i++ {
		if b.min[i] > r.max[i] || b.max[i] < r.min[i] {
			return false
		}
	}
	return true
}

func (r *rRect) search(
	target *rRect, height int,
	iter func(min, max []float64, value interface{}) bool,
) bool {
	n := r.data.(*rNode)
	if height == 0 {
		for i := 0; i < n.count; i++ {
			if target.intersects(&n.rects[i]) {
				if !iter(n.rects[i].min[:], n.rects[i].max[:],
					n.rects[i].data) {
					return false
				}
			}
		}
	} else {
		for i := 0; i < n.count; i++ {
			if target.intersects(&n.rects[i]) {
				if !n.rects[i].search(target, height-1, iter) {
					return false
				}
			}
		}
	}
	return true
}

func (tr *rTree) search(
	target *rRect,
	iter func(min, max []float64, value interface{}) bool,
) {
	if tr.root.data == nil {
		return
	}
	if target.intersects(&tr.root) {
		tr.root.search(target, tr.height, iter)
	}
}

// Search ...
func (tr *rTree) Search(
	min, max []float64,
	iter func(min, max []float64, value interface{}) bool,
) {
	var target rRect
	fit(min, max, nil, &target)
	tr.search(&target, iter)
}

func appendFloat(dst []byte, num float64) []byte {
	var buf [8]byte
	binary.LittleEndian.PutUint64(buf[:], math.Float64bits(num))
	return append(dst, buf[:]...)
}

func (r *rRect) compress(dst []byte, height int) []byte {
	n := r.data.(*rNode)
	dst = appendFloat(dst, r.min[0])
	dst = appendFloat(dst, r.min[1])
	dst = appendFloat(dst, r.max[0])
	dst = appendFloat(dst, r.max[1])
	dst = append(dst, byte(n.count))
	if height == 0 {
		var ibytes byte = 1
		for i := 0; i < n.count; i++ {
			ibytes2 := numBytes(uint32(n.rects[i].data.(int)))
			if ibytes2 > ibytes {
				ibytes = ibytes2
			}
		}
		dst = append(dst, ibytes)
		for i := 0; i < n.count; i++ {
			dst = appendNum(dst, uint32(n.rects[i].data.(int)), ibytes)
		}
		return dst
	}
	mark := make([]int, n.count)
	for i := 0; i < n.count; i++ {
		mark[i] = len(dst)
		dst = append(dst, 0, 0, 0, 0)
	}
	for i := 0; i < n.count; i++ {
		binary.LittleEndian.PutUint32(dst[mark[i]:], uint32(len(dst)))
		dst = n.rects[i].compress(dst, height-1)
	}
	return dst
}

func (tr *rTree) compress(dst []byte) []byte {
	if tr.root.data == nil {
		return dst
	}
	dst = append(dst, byte(tr.height))
	return tr.root.compress(dst, tr.height)
}

func rCompressSearch(
	data []byte,
	addr int,
	series *baseSeries,
	rect Rect,
	iter func(seg Segment, item int) bool,
) bool {
	if int(addr) == len(data) {
		return true
	}
	height := int(data[addr])
	addr++
	return rnCompressSearch(data, addr, series, rect, height, iter)
}

func rnCompressSearch(
	data []byte,
	addr int,
	series *baseSeries,
	rect Rect,
	height int,
	iter func(seg Segment, item int) bool,
) bool {
	var nrect Rect
	nrect.Min.X = math.Float64frombits(binary.LittleEndian.Uint64(data[addr:]))
	addr += 8
	nrect.Min.Y = math.Float64frombits(binary.LittleEndian.Uint64(data[addr:]))
	addr += 8
	nrect.Max.X = math.Float64frombits(binary.LittleEndian.Uint64(data[addr:]))
	addr += 8
	nrect.Max.Y = math.Float64frombits(binary.LittleEndian.Uint64(data[addr:]))
	addr += 8
	if !rect.IntersectsRect(nrect) {
		return true
	}
	count := int(data[addr])
	addr++
	if height == 0 {
		ibytes := data[addr]
		addr++
		for i := 0; i < count; i++ {
			item := int(readNum(data[addr:], ibytes))
			addr += int(ibytes)
			seg := series.SegmentAt(int(item))
			irect := seg.Rect()
			if irect.IntersectsRect(rect) {
				if !iter(seg, int(item)) {
					return false
				}
			}
		}
		return true
	}
	for i := 0; i < count; i++ {
		naddr := int(binary.LittleEndian.Uint32(data[addr:]))
		addr += 4
		if !rnCompressSearch(data, naddr, series, rect, height-1, iter) {
			return false
		}
	}
	return true
}