EngoEngine / engo

Engo is an open-source 2D game engine written in Go.
https://engoengine.github.io
MIT License
1.75k stars 136 forks source link

quadtree to fast #804

Open zswDev opened 6 months ago

zswDev commented 6 months ago

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