quadtree.go

package lib

import (
	"math"
)

const (
	QuadTreeCapacity = 32
	QuadTreeMaxDepth = 12
)

type Vector2Int struct {
	X int32
	Y int32
}

func (v Vector2Int) Add(v2 Vector2Int) Vector2Int {
	v.X += v2.X
	v.Y += v2.Y
	return v
}

func (v Vector2Int) Sub(v2 Vector2Int) Vector2Int {
	v.X -= v2.X
	v.Y -= v2.Y
	return v
}

func (v Vector2Int) Divide(divisor int32) Vector2Int {
	v.X /= divisor
	v.Y /= divisor
	return v
}

func (v Vector2Int) Multiply(scalar int32) Vector2Int {
	v.X *= scalar
	v.Y *= scalar
	return v
}

func (v Vector2Int) Mod(divisor int32) Vector2Int {
	v.X %= divisor
	v.Y %= divisor
	return v
}

func (v Vector2Int) ToVector2() Vector2 {
	return Vector2{float64(v.X), float64(v.Y)}
}

type Vector2 struct {
	X float64
	Y float64
}

func (v Vector2) Add(v2 Vector2) Vector2 {
	v.X += v2.X
	v.Y += v2.Y
	return v
}

func (v Vector2) Sub(v2 Vector2) Vector2 {
	v.X -= v2.X
	v.Y -= v2.Y
	return v
}

func (v Vector2) Multiply(scalar float64) Vector2 {
	v.X *= scalar
	v.Y *= scalar
	return v
}

func (v Vector2) Magnitude() float64 {
	return math.Sqrt(v.X*v.X + v.Y*v.Y)
}

func (v Vector2) Normalize() Vector2 {
	div := v.Magnitude()
	v.X /= div
	v.Y /= div
	return v
}

func (v Vector2) Floor() Vector2 {
	v.X = float64(int(v.X))
	v.Y = float64(int(v.Y))
	return v
}

func (v Vector2) ToVector2Int() Vector2Int {
	return Vector2Int{int32(v.X), int32(v.Y)}
}

type Rectangle struct {
	X      int32
	Y      int32
	Width  int32
	Height int32
}

func (r Rectangle) Contains(v Vector2Int) bool {
	return r.X <= v.X && r.Y <= v.Y && r.X+r.Width > v.X && r.Y+r.Height > v.Y
}

func (r Rectangle) ContainsRect(r2 Rectangle) bool {
	return r.X <= r2.X && r.X+r.Width >= r2.X+r2.Width && r.Y <= r2.Y && r.Y+r.Height >= r2.Y+r2.Height
}

func (r Rectangle) Intersects(r2 Rectangle) bool {
	return r.X < r2.X+r2.Width && r.X+r.Width > r2.X && r.Y < r2.Y+r2.Height && r.Y+r.Height > r2.Y
}

func (r Rectangle) Move(v Vector2Int) Rectangle {
	return Rectangle{
		X:      r.X + v.X,
		Y:      r.Y + v.Y,
		Width:  r.Width,
		Height: r.Height,
	}
}

func (r Rectangle) Grow(v Vector2Int) Rectangle {
	return Rectangle{
		X:      r.X,
		Y:      r.Y,
		Width:  r.Width + v.X,
		Height: r.Height + v.Y,
	}
}

func (r Rectangle) Merge(r2 Rectangle) (Rectangle, bool) {
	if r.ContainsRect(r2) {
		return r, true
	}
	if r2.ContainsRect(r) {
		return r2, true
	}

	// is it vertically aligned?
	if r.X == r2.X && r.Width == r2.Width {
		// if the first rect is above the second
		if r.Y < r2.Y {
			// if the first rect touches the second
			if r.Y+r.Height >= r2.Y {
				r.Height = r2.Y + r2.Height - r.Y
				return r, true
			}
			return r, false
		} else {
			if r2.Y+r2.Height >= r.Y {
				r2.Height = r.Y + r.Height - r2.Y
				return r2, true
			}
			return r, false
		}
	}

	// is it horizontally aligned?
	if r.Y == r2.Y && r.Height == r2.Height {
		// if the first rect is left of the second
		if r.X < r2.X {
			// if the first rect touches the second
			if r.X+r.Width >= r2.X {
				r.Width = r2.X + r2.Width - r.X
				return r, true
			}
			return r, false
		} else {
			if r2.X+r2.Width >= r.X {
				r2.Width = r.X + r.Width - r2.X
				return r2, true
			}
			return r, false
		}
	}

	return r, false
}

type QuadTreeEntry struct {
	ID       uint64
	Position Rectangle
	next     int
}

type quadTreeNode struct {
	firstChild int
	count      int
}

// QuadTree is a partitioned space structure for storing rectangles with an associated ID.
// It is used to efficiently search regions of space for elements within.
type QuadTree struct {
	bounds  Rectangle
	nodes   FreeList[quadTreeNode]
	entries FreeList[QuadTreeEntry]
}

func NewQuadTree(bounds Rectangle) *QuadTree {
	qt := &QuadTree{
		bounds:  bounds,
		nodes:   FreeList[quadTreeNode]{firstFree: -1},
		entries: FreeList[QuadTreeEntry]{firstFree: -1},
	}

	qt.nodes.Insert(quadTreeNode{firstChild: -1})

	return qt
}

func (q *QuadTree) Insert(e QuadTreeEntry) {
	e.next = -1
	index := q.entries.Insert(e)
	q.insert(index, e.Position, q.bounds, 0, 0)
}

func (q *QuadTree) insert(index int, position, bounds Rectangle, depth, nodeIndex int) {
	// if the player does not exist within our bounds, do nothing.
	// one of our siblings will accept the player
	if !bounds.Intersects(position) {
		return
	}

	// get the node
	node := q.nodes.Get(nodeIndex)

	// if this is a branch
	if q.isBranchNode(node) {
		w := bounds.Width / 2
		h := bounds.Height / 2
		q.insert(index, position, Rectangle{bounds.X + w, bounds.Y + h, w, h}, depth+1, node.firstChild)
		q.insert(index, position, Rectangle{bounds.X, bounds.Y + h, w, h}, depth+1, node.firstChild+1)
		q.insert(index, position, Rectangle{bounds.X, bounds.Y, w, h}, depth+1, node.firstChild+2)
		q.insert(index, position, Rectangle{bounds.X + w, bounds.Y, w, h}, depth+1, node.firstChild+3)
	} else {
		// put element at start of Raw linked list
		entry := q.entries.Get(index)
		entry.next = node.firstChild
		q.entries.Set(index, entry)
		node.firstChild = index
		node.count++

		// split
		if node.count > QuadTreeCapacity && depth < QuadTreeMaxDepth {
			// save the index of the first Raw element
			currentChild := node.firstChild

			// create a new 4 quad trees and set the first as first index
			node.firstChild = q.nodes.Insert(quadTreeNode{firstChild: -1})
			q.nodes.Insert(quadTreeNode{firstChild: -1})
			q.nodes.Insert(quadTreeNode{firstChild: -1})
			q.nodes.Insert(quadTreeNode{firstChild: -1})

			// split out children into leaf nodes
			w := bounds.Width / 2
			h := bounds.Height / 2

			// insert this nodes old elements
			for {
				if currentChild == -1 {
					break
				}
				entry := q.entries.Get(currentChild)
				q.insert(currentChild, entry.Position, Rectangle{bounds.X + w, bounds.Y + h, w, h}, depth+1, node.firstChild)
				q.insert(currentChild, entry.Position, Rectangle{bounds.X, bounds.Y + h, w, h}, depth+1, node.firstChild+1)
				q.insert(currentChild, entry.Position, Rectangle{bounds.X, bounds.Y, w, h}, depth+1, node.firstChild+2)
				q.insert(currentChild, entry.Position, Rectangle{bounds.X + w, bounds.Y, w, h}, depth+1, node.firstChild+3)
				currentChild = entry.next
			}
			node.count = -1
		}

		q.nodes.Set(nodeIndex, node)
	}

	return
}

func (q *QuadTree) Remove(e QuadTreeEntry) {
	q.remove(e, q.bounds, 0)
}

func (q *QuadTree) remove(e QuadTreeEntry, bounds Rectangle, nodeIndex int) {
	// if the entity does not exist within our bounds, do nothing.
	// the entity could never have been in our tree.
	if !bounds.Intersects(e.Position) {
		return
	}

	// get the node
	node := q.nodes.Get(nodeIndex)

	// find and remove item
	if q.isBranchNode(node) {
		// ask children to remove it
		w := bounds.Width / 2
		h := bounds.Height / 2

		q.remove(e, Rectangle{bounds.X + w, bounds.Y + h, w, h}, node.firstChild)
		q.remove(e, Rectangle{bounds.X, bounds.Y + h, w, h}, node.firstChild+1)
		q.remove(e, Rectangle{bounds.X, bounds.Y, w, h}, node.firstChild+2)
		q.remove(e, Rectangle{bounds.X + w, bounds.Y, w, h}, node.firstChild+3)
	} else {
		currentChild := node.firstChild
		lastCheckedIndex := -1
		for {
			if currentChild == -1 {
				break
			}

			// get the child element we're currently checking
			entry := q.entries.Get(currentChild)
			// if the child element has a matching id
			if entry.ID == e.ID {
				// free the element from the list
				q.entries.Erase(currentChild)
				node.count--
				// if the element we removed was the first in the list
				if lastCheckedIndex == -1 {
					// set the new first child to be the second item in the list
					node.firstChild = entry.next
				} else { // otherwise if there was an element before the one that was removed
					// get the last element and update it to point to the removed elements next ptr
					prevEntry := q.entries.Get(lastCheckedIndex)
					prevEntry.next = entry.next
					q.entries.Set(lastCheckedIndex, prevEntry)
				}
				q.nodes.Set(nodeIndex, node)
				return
			}
			// save the index of the last checked element
			lastCheckedIndex = currentChild
			// set the next child to check
			currentChild = entry.next
		}
	}
}

func (q *QuadTree) EntriesWithin(results *[]QuadTreeEntry, rect Rectangle) {
	q.entriesWithin(results, rect, q.bounds, 0)
}

func (q *QuadTree) entriesWithin(results *[]QuadTreeEntry, rect, bounds Rectangle, nodeIndex int) {
	if !bounds.Intersects(rect) {
		return
	}

	node := q.nodes.Get(nodeIndex)

	// find and remove item
	if q.isBranchNode(node) {
		// ask children to search
		w := bounds.Width / 2
		h := bounds.Height / 2

		q.entriesWithin(results, rect, Rectangle{bounds.X + w, bounds.Y + h, w, h}, node.firstChild)
		q.entriesWithin(results, rect, Rectangle{bounds.X, bounds.Y + h, w, h}, node.firstChild+1)
		q.entriesWithin(results, rect, Rectangle{bounds.X, bounds.Y, w, h}, node.firstChild+2)
		q.entriesWithin(results, rect, Rectangle{bounds.X + w, bounds.Y, w, h}, node.firstChild+3)
	} else {
		currentChild := node.firstChild
		for {
			if currentChild == -1 {
				break
			}

			// get the child element we're currently checking
			entry := q.entries.Get(currentChild)
			// if the child element is in the search rectangle
			if rect.Intersects(entry.Position) {
				*results = append(*results, entry)
			}
			currentChild = entry.next
		}
	}
}

func (q *QuadTree) CleanUp() {
	var stack []int

	// Only process the root if it's a branch
	if q.isBranchNode(q.nodes.Get(0)) {
		stack = append(stack, 0)
	}

	for {
		if len(stack) == 0 {
			break
		}

		// pop stack
		n := len(stack) - 1
		nodeToProcess := stack[n]
		stack = stack[:n]

		node := q.nodes.Get(nodeToProcess)

		// Loop through the children.
		var numberOfEmptyLeaves int
		for i := 0; i < 4; i++ {
			childIndex := node.firstChild + i
			childNode := q.nodes.Get(childIndex)

			// Increment empty leaf count if the child is an empty
			// leaf. Otherwise if the child is a branch, add it to
			// the stack to be processed in the next iteration.
			if childNode.count == 0 {
				numberOfEmptyLeaves++
			} else if q.isBranchNode(childNode) {
				stack = append(stack, childIndex)
			}
		}

		// If all the children were empty leaves, remove them and
		// make this node the new empty leaf.
		if numberOfEmptyLeaves == 4 {
			// Push all 4 children to the free list.
			for i := 3; i >= 0; i-- {
				q.nodes.Erase(node.firstChild + i)
			}
			// Make this node the new empty leaf.
			node.firstChild = -1
			node.count = 0
			q.nodes.Set(nodeToProcess, node)
		}
	}
}

func (q *QuadTree) Clear() {
	q.nodes.Clear()
	q.entries.Clear()
	q.nodes.Insert(quadTreeNode{firstChild: -1})
}

func (q *QuadTree) isBranchNode(node quadTreeNode) bool {
	return node.count == -1
}

func (q *QuadTree) Walk(f func(i int, r Rectangle, entries []QuadTreeEntry)) {
	q.walk(f, 0, q.bounds, 0)
}

func (q *QuadTree) walk(f func(i int, r Rectangle, entries []QuadTreeEntry), quadrant int, bounds Rectangle, nodeIndex int) {
	node := q.nodes.Get(nodeIndex)

	// find and remove item
	if q.isBranchNode(node) {
		// ask children to search
		w := bounds.Width / 2
		h := bounds.Height / 2

		q.walk(f, 0, Rectangle{bounds.X + w, bounds.Y + h, w, h}, node.firstChild)
		q.walk(f, 1, Rectangle{bounds.X, bounds.Y + h, w, h}, node.firstChild+1)
		q.walk(f, 2, Rectangle{bounds.X, bounds.Y, w, h}, node.firstChild+2)
		q.walk(f, 3, Rectangle{bounds.X + w, bounds.Y, w, h}, node.firstChild+3)
	} else {
		var entries []QuadTreeEntry
		currentChild := node.firstChild
		for {
			if currentChild == -1 {
				break
			}

			// get the child element we're currently checking
			entry := q.entries.Get(currentChild)
			// if the child element is in the search rectangle
			entries = append(entries, entry)
			currentChild = entry.next
		}
		f(quadrant, bounds, entries)
	}
}