fluids/quadtree/quadtree.go

package quadtree

import (
	"fluids/colliders"
	"fluids/gamedata"
)

const (
	QuadtreeMaxColliders = 40
)

type Quadrant struct {
	Position   gamedata.Vector
	Dimensions gamedata.Vector
}

type Quadtree struct {
	quadrant  Quadrant
	children  []*Quadtree
	colliders []colliders.Collider
}

func New(quadrant Quadrant) *Quadtree {
	qt := &Quadtree{
		quadrant: quadrant,
	}
	return qt
}

func (q *Quadtree) Insert(obj colliders.Collider) bool {

	var result bool = false
	//check that object meets containment criteria
	if !q.ContainsPoint(obj.GetPosition()) {
		return result
	}

	//if we have children, we go deeper
	if len(q.children) > 0 {
		for _, node := range q.children {
			if node.Insert(obj) {
				result = true
				break
			}
		}
	} else if len(q.colliders) < QuadtreeMaxColliders {
		//otherwise, if we have space, add
		q.colliders = append(q.colliders, obj)
		result = true
	} else {
		//otherwise we have to subdivide, reorganize, and add
		result = q.SubdivideAndInsert(obj)
	}

	return result
}

func (q *Quadtree) ContainsPoint(p gamedata.Vector) bool {
	var result bool = false
	if (q.quadrant.Position.X-q.quadrant.Dimensions.X/2) <= p.X &&
		p.X <= (q.quadrant.Position.X+q.quadrant.Dimensions.X/2) &&
		(q.quadrant.Position.Y-q.quadrant.Dimensions.Y/2) <= p.Y &&
		p.Y <= (q.quadrant.Position.Y+q.quadrant.Dimensions.Y/2) {
		result = true
	}
	return result
}

func (q *Quadtree) Remove(obj colliders.Collider) bool {
	var result bool = false

	if !q.ContainsPoint(obj.GetPosition()) {
		return result
	}

	if len(q.colliders) > 0 {
		//examine existing colliders
		var collection []colliders.Collider
		for _, node := range q.colliders {
			if node == obj {
				result = true
			} else {
				collection = append(collection, node)
			}
		}
		q.colliders = collection
	} else {
		//need to check children
		for _, child := range q.children {
			if child.Remove(obj) {
				result = true
				break
			}
		}
	}

	return result
}

func (q *Quadtree) SubdivideAndInsert(obj colliders.Collider) bool {

	var result bool = false

	//initialize up children
	q.children = q.children[:0]
	q.children = append(q.children, New(Quadrant{Position: gamedata.Vector{X: q.quadrant.Position.X - q.quadrant.Dimensions.X/4, Y: q.quadrant.Position.Y - q.quadrant.Dimensions.Y/4}, Dimensions: gamedata.Vector{X: q.quadrant.Dimensions.X / 2, Y: q.quadrant.Dimensions.Y / 2}}))
	q.children = append(q.children, New(Quadrant{Position: gamedata.Vector{X: q.quadrant.Position.X + q.quadrant.Dimensions.X/4, Y: q.quadrant.Position.Y - q.quadrant.Dimensions.Y/4}, Dimensions: gamedata.Vector{X: q.quadrant.Dimensions.X / 2, Y: q.quadrant.Dimensions.Y / 2}}))
	q.children = append(q.children, New(Quadrant{Position: gamedata.Vector{X: q.quadrant.Position.X - q.quadrant.Dimensions.X/4, Y: q.quadrant.Position.Y + q.quadrant.Dimensions.Y/4}, Dimensions: gamedata.Vector{X: q.quadrant.Dimensions.X / 2, Y: q.quadrant.Dimensions.Y / 2}}))
	q.children = append(q.children, New(Quadrant{Position: gamedata.Vector{X: q.quadrant.Position.X + q.quadrant.Dimensions.X/4, Y: q.quadrant.Position.Y + q.quadrant.Dimensions.Y/4}, Dimensions: gamedata.Vector{X: q.quadrant.Dimensions.X / 2, Y: q.quadrant.Dimensions.Y / 2}}))

	//move colliders into child nodes
	var failed bool = false
	for _, collider := range q.colliders {
		var added bool = false
		for _, node := range q.children {
			if node.Insert(collider) {
				added = true
				break
			}
		}

		//couldn't add collider into any subdivided node; failure
		if !added {
			failed = true
			break
		}
	}

	//if all went well:
	// a. need to wipe colliders from current node
	// b. need to now insert new collider into one of the children
	if !failed {
		q.colliders = q.colliders[:0]

		for _, node := range q.children {
			if node.Insert(obj) {
				result = true
				break
			}
		}
	}

	return result
}

func (q *Quadtree) Clear() {
	q.colliders = q.colliders[:0]
	q.children = q.children[:0]
}

func (q *Quadtree) GetQuadrants() []Quadrant {
	var result []Quadrant

	//if we have children, traverse and add their quadrants
	if len(q.children) > 0 {
		for _, child := range q.children {
			result = append(result, child.GetQuadrants()...)
		}
	}

	result = append(result, q.quadrant)

	return result
}

// find and retrieve all colliders in given quadrant
func (q *Quadtree) FindAll(quadrant Quadrant) []colliders.Collider {
	var result []colliders.Collider

	//search coordinates
	sx0 := quadrant.Position.X - quadrant.Dimensions.X/2
	sx1 := quadrant.Position.X + quadrant.Dimensions.X/2
	sy0 := quadrant.Position.Y - quadrant.Dimensions.Y/2
	sy1 := quadrant.Position.Y + quadrant.Dimensions.Y/2

	//quadtree coordinates
	qx0 := q.quadrant.Position.X - q.quadrant.Dimensions.X/2
	qx1 := q.quadrant.Position.X + q.quadrant.Dimensions.X/2
	qy0 := q.quadrant.Position.Y - q.quadrant.Dimensions.Y/2
	qy1 := q.quadrant.Position.Y + q.quadrant.Dimensions.Y/2

	//AABB check
	if sx0 < qx1 && sx1 > qx0 &&
		sy0 < qy1 && sy1 > qy0 {

		//if we have children, check each of those
		for _, node := range q.children {
			result = append(result, node.FindAll(quadrant)...)
		}

		result = append(result, q.colliders...)
	}

	return result
}