quadtree to fast

[zswDev]

Reduce memory allocation

package fast

import (
	"room/business/tool"
	"sync"
)

//https://github.com/EngoEngine/engo

var (
	quadtreeNodePool *sync.Pool
)

const minQuadtreeCellSize = 0.01

func init() {
	quadtreeNodePool = &sync.Pool{
		New: func() interface{} {
			return &quadtreeNode{
				Objects: []quadtreeNodeData{},
			}
		},
	}
}

type AABB = tool.AABB

type Point = tool.Point

type AABBer = tool.AABBer

// type AABB struct {
// 	Min, Max Point
// }

// // AABBer is an interface for everything that provides information about its axis aligned bounding box.
// type AABBer interface {
// 	// AABB returns the axis aligned bounding box.
// 	AABB() AABB
// }

func aabbOverlaps(a, b AABB) bool {
	// a is left of b
	if a.Max.X < b.Min.X {
		return false
	}

	// a is right of b
	if a.Min.X > b.Max.X {
		return false
	}

	// a is above b
	if a.Max.Y < b.Min.Y {
		return false
	}

	// a is below b
	if a.Min.Y > b.Max.Y {
		return false
	}

	// The two overlap
	return true
}

func aabbWidth(x AABB) float64 {
	return x.Max.X - x.Min.X
}

func aabbHeight(x AABB) float64 {
	return x.Max.Y - x.Min.Y
}

func aabbRect(x, y, width, height float64) AABB {
	return AABB{
		Min: Point{
			X: x,
			Y: y,
		},
		Max: Point{
			X: x + width,
			Y: y + height,
		},
	}
}

type quadtreeNodeData struct {
	Value AABBer
	AABB  AABB
}

type quadtreeNode struct {
	hasNodes bool
	Level    int
	Bounds   AABB
	Tree     *Quadtree
	Nodes    [4]*quadtreeNode
	Objects  []quadtreeNodeData
}

// Quadtree implementation which can store AABBer values
type Quadtree struct {
	MaxObjects int // Maximum objects a node can hold before splitting into 4 subnodes
	MaxLevels  int // Total max levels inside root Quadtree
	Total      int
	root       *quadtreeNode
	result     []AABBer
}

func calcMaxLevel(width, height float64) int {
	res := 0
	for width > minQuadtreeCellSize && height > minQuadtreeCellSize {
		res++
		width, height = width/2, height/2
	}
	return res
}

// NewQuadtree creates a new quadtree for the given bounds.
// When setting usePool to true, the internal values will be taken from a sync.Pool which reduces the allocation overhead.
// maxObjects tells the tree how many objects should be stored within a level before the quadtree cell is split.
func NewQuadtree(bounds AABB, maxObjects int) *Quadtree {
	qt := &Quadtree{
		MaxObjects: maxObjects,
		result:     []tool.AABBer{},
	}
	qt.root = qt.newNode(bounds, 0)
	qt.MaxLevels = calcMaxLevel(aabbWidth(bounds), aabbHeight(bounds))
	return qt
}

// Destroy frees the nodes if the Quadtree uses the node pool
func (qt *Quadtree) Destroy() {
	qt.freeQuadtreeNode(qt.root)
	qt.root = nil
}

func (qt *Quadtree) newNode(bounds AABB, level int) (node *quadtreeNode) {

	node = quadtreeNodePool.Get().(*quadtreeNode)
	if node == nil {
		node = &quadtreeNode{
			Objects: []quadtreeNodeData{},
		}
	}
	node.Tree = qt
	node.Bounds = bounds
	node.Level = level
	return node
}

func (qt *Quadtree) freeQuadtreeNode(n *quadtreeNode) {

	if n.hasNodes {
		for i := range n.Nodes {
			qt.freeQuadtreeNode(n.Nodes[i])
			n.Nodes[i] = nil
		}
	}
	n.Objects = n.Objects[:0]
	n.Tree = nil
	n.hasNodes = false
	quadtreeNodePool.Put(n)
}

// split - split the node into 4 subnodes
func (qt *quadtreeNode) split() {
	if qt.hasNodes {
		return
	}
	qt.hasNodes = true

	nextLevel := qt.Level + 1
	subWidth := aabbWidth(qt.Bounds) / 2
	subHeight := aabbHeight(qt.Bounds) / 2
	x := qt.Bounds.Min.X
	y := qt.Bounds.Min.Y

	//top right node (0)
	qt.Nodes[0] = qt.Tree.newNode(aabbRect(x+subWidth, y, subWidth, subHeight), nextLevel)

	//top left node (1)
	qt.Nodes[1] = qt.Tree.newNode(aabbRect(x, y, subWidth, subHeight), nextLevel)

	//bottom left node (2)
	qt.Nodes[2] = qt.Tree.newNode(aabbRect(x, y+subHeight, subWidth, subHeight), nextLevel)

	//bottom right node (3)
	qt.Nodes[3] = qt.Tree.newNode(aabbRect(x+subWidth, y+subHeight, subWidth, subHeight), nextLevel)
}

func (qt *quadtreeNode) isEmpty() bool {
	return len(qt.Objects) == 0 && !qt.hasNodes
}

func (qt *quadtreeNode) unsplit() {
	for i := 0; i < 4; i++ {
		if !qt.Nodes[i].isEmpty() {
			return
		}
	}
	for i := 0; i < 4; i++ {
		qt.Tree.freeQuadtreeNode(qt.Nodes[i])
		qt.Nodes[i] = nil
	}
	qt.hasNodes = false
}

// getIndex - Determine which quadrant the object belongs to (0-3)
func (qt *quadtreeNode) getIndex(pRect AABB) int {
	horzMidpoint := qt.Bounds.Min.X + (aabbWidth(qt.Bounds) / 2)
	vertMidpoint := qt.Bounds.Min.Y + (aabbHeight(qt.Bounds) / 2)

	//pRect can completely fit within the top quadrants
	topQuadrant := (pRect.Min.Y < vertMidpoint) && (pRect.Max.Y < vertMidpoint)

	//pRect can completely fit within the bottom quadrants
	bottomQuadrant := (pRect.Min.Y > vertMidpoint)

	//pRect can completely fit within the left quadrants
	if (pRect.Min.X < horzMidpoint) && (pRect.Max.X < horzMidpoint) {
		if topQuadrant {
			return 1
		} else if bottomQuadrant {
			return 2
		}
	} else if pRect.Min.X > horzMidpoint {
		//pRect can completely fit within the right quadrants
		if topQuadrant {
			return 0
		} else if bottomQuadrant {
			return 3
		}
	}

	return -1 // index of the subnode (0-3), or -1 if pRect cannot completely fit within a subnode and is part of the parent node
}

// Insert inserts the given item to the quadtree
func (qt *Quadtree) Insert(item AABBer) {
	qt.Total++
	qt.root.Insert(quadtreeNodeData{AABB: item.AABB(), Value: item})
}

func (qt *quadtreeNode) Insert(item quadtreeNodeData) {
	if qt.hasNodes {
		index := qt.getIndex(item.AABB)
		if index != -1 {
			qt.Nodes[index].Insert(item)
			return
		}
	}

	// If we don't subnodes within the Quadtree
	qt.Objects = append(qt.Objects, item)

	// If total objects is greater than max objects and level is less than max levels
	if (len(qt.Objects) > qt.Tree.MaxObjects) && (qt.Tree.MaxLevels <= 0 || qt.Level < qt.Tree.MaxLevels) {
		// split if we don't already have subnodes
		if !qt.hasNodes {
			qt.split()
		}

		// Add all objects to there corresponding subNodes
		olen := len(qt.Objects)
		for i := 0; i < len(qt.Objects); {
			object := qt.Objects[i] // Get the object out of the slice
			bounds := object.AABB
			index := qt.getIndex(bounds)
			if index != -1 {
				qt.Objects = removeObjects(qt.Objects, i, olen)
				olen--
				qt.Nodes[index].Insert(object)
			} else {
				i++
			}
		}
	}
}
func removeObjects(objects []quadtreeNodeData, delIdx, olen int) []quadtreeNodeData {
	endIdx := olen - 1
	if endIdx != delIdx {
		objects[delIdx] = objects[endIdx]
	}
	return objects[:endIdx]
}

func (qt *quadtreeNode) Remove(item AABBer, pRect AABB) {
	if qt.hasNodes {
		index := qt.getIndex(pRect)
		if index != -1 {
			qt.Nodes[index].Remove(item, pRect)
			qt.unsplit()
			return
		}
	}
	olen := len(qt.Objects)
	for i := 0; i < len(qt.Objects); i++ {
		if qt.Objects[i].Value == item {
			removeObjects(qt.Objects, i, olen)
			return
		}
	}
}

// Remove removes the given item from the quadtree
func (qt *Quadtree) Remove(item AABBer) {
	bounds := item.AABB()
	qt.root.Remove(item, bounds)
}

// Retrieve returns all objects that could collide with the given bounding box
func (qt *quadtreeNode) Retrieve(pRect AABB, result []tool.AABBer, filter func(aabb AABBer) bool) []AABBer {

	index := qt.getIndex(pRect)

	// Array with all detected objects
	for i := range qt.Objects {
		if aabbOverlaps(pRect, qt.Objects[i].Value.AABB()) &&
			(filter == nil || filter(qt.Objects[i].Value)) {

			result = append(result, qt.Objects[i].Value)
		}
	}

	//if we have subnodes ...
	if qt.hasNodes {
		//if pRect fits into a subnode ..
		if index != -1 {
			result = qt.Nodes[index].Retrieve(pRect, result, filter)
		} else {
			//if pRect does not fit into a subnode, check it against all subnodes
			for i := 0; i < 4; i++ {
				result = qt.Nodes[i].Retrieve(pRect, result, filter)
			}
		}
	}

	return result

}

// Retrieve returns all objects that could collide with the given bounding box and passing the given filter function.
func (qt *Quadtree) Retrieve(find AABB, filter func(aabb AABBer) bool) []AABBer {

	qt.result = qt.result[:0]
	potentials := qt.root.Retrieve(find, qt.result, filter)
	return potentials
}

// Clear removes all items from the quadtree
func (qt *Quadtree) Clear() {
	bounds := qt.root.Bounds
	qt.freeQuadtreeNode(qt.root)
	qt.root = qt.newNode(bounds, 0)
	qt.Total = 0
	qt.result = qt.result[:0]
}