aoi/quadtree.go

package aoi

import "sync"

const (
	leftUp int = iota
	rightUp
	leftDown
	rightDown

	maxCap  = 500 // 节点最大容量
	maxDeep = 4   // 节点最大深度
	radius  = 16  // 视野半径
)

type Node struct {
	Leaf      bool      // 是否为叶子节点
	Deep      int       // 深度
	AreaWidth float64   // 格子宽度(长=宽)
	XStart    float64   // 起始范围
	YStart    float64   // 起始范围
	Tree      *QuadTree // 树指针
	Child     [4]*Node  // 子节点
	Entities  *sync.Map // 实体
}

type QuadTree struct {
	maxCap, maxDeep int
	radius          float64
	mPool           sync.Pool
	*Node
}

func NewSonNode(xStart, yStart float64, parent *Node) *Node {
	son := &Node{
		Leaf:      true,
		Deep:      parent.Deep + 1,
		AreaWidth: parent.AreaWidth / 2,
		XStart:    xStart,
		YStart:    yStart,
		Tree:      parent.Tree,
		Entities:  parent.Tree.mPool.Get().(*sync.Map),
	}

	return son
}

// canCut 检查节点是否可以分割
func (n *Node) canCut() bool {
	if n.XStart+n.AreaWidth/2 > 0 && n.YStart+n.AreaWidth/2 > 0 {
		return true
	}
	return false
}

// needCut 检查节点是否需要分割
func (n *Node) needCut() bool {
	lens := 0
	n.Entities.Range(func(key, value interface{}) bool {
		lens++
		return true
	})
	return lens+1 > n.Tree.maxCap && n.Deep+1 <= n.Tree.maxDeep && n.canCut()
}

// intersects 检查坐标是否在节点范围内
func (n *Node) intersects(x, y float64) bool {
	if n.XStart <= x && x < n.XStart+n.AreaWidth && n.YStart <= y && y < n.YStart+n.AreaWidth {
		return true
	}
	return false
}

// findSonQuadrant 根据坐标寻找子节点的方位
func (n *Node) findSonQuadrant(x, y float64) int {
	if x < n.Child[rightDown].XStart {
		if y < n.Child[rightDown].YStart {
			return leftUp
		}
		return leftDown
	}
	if y < n.Child[rightDown].YStart {
		return rightUp
	}
	return rightDown
}

// cutNode 分割节点
func (n *Node) cutNode() {
	n.Leaf = false
	half := n.AreaWidth / 2

	n.Child[leftUp] = NewSonNode(n.XStart, n.YStart, n)
	n.Child[rightUp] = NewSonNode(n.XStart+half, n.YStart, n)
	n.Child[leftDown] = NewSonNode(n.XStart, n.YStart+half, n)
	n.Child[rightDown] = NewSonNode(n.XStart+half, n.YStart+half, n)

	// 将实体迁移到对应子节点
	n.Entities.Range(func(k, v interface{}) bool {
		entity := v.(*Entity)
		for _, node := range n.Child {
			if node.intersects(entity.X, entity.Y) {
				node.Entities.Store(entity.Key, entity)
			}
		}
		n.Entities.Delete(k)
		return true
	})

	n.Tree.mPool.Put(n.Entities)
	n.Entities = nil
}

func NewQuadTree(xStart, yStart, width float64) AOI {
	basicNode := &Node{
		Leaf:      true,
		Deep:      1,
		AreaWidth: width,
		XStart:    xStart,
		YStart:    yStart,
		Child:     [4]*Node{},
	}
	tree := &QuadTree{
		maxDeep: maxDeep,
		maxCap:  maxCap,
		radius:  radius,
		Node:    basicNode,
	}
	tree.mPool.New = func() interface{} {
		return &sync.Map{}
	}
	basicNode.Tree = tree
	basicNode.Entities = tree.mPool.Get().(*sync.Map)
	return tree
}

func (n *Node) Add(x, y float64, name string) {
	// 判断是否需要分割
	if n.Leaf && n.needCut() {
		n.cutNode()
	}

	// 非叶子节点往下递归
	if !n.Leaf {
		n.Child[n.findSonQuadrant(x, y)].Add(x, y, name)
		return
	}

	entity := entityPool.Get().(*Entity)
	entity.X = x
	entity.Y = y
	entity.Key = name

	// 叶子节点进行存储
	n.Entities.Store(entity.Key, entity)
}

func (n *Node) Delete(x, y float64, name string) {
	if !n.Leaf {
		n.Child[n.findSonQuadrant(x, y)].Delete(x, y, name)
		return
	}

	if entity, ok := n.Entities.Load(name); ok {
		n.Entities.Delete(name)
		entityPool.Put(entity)
	}
}

func (n *Node) Search(x, y float64) []string {
	result := resultPool.Get().([]string)
	defer func() {
		result = result[:0]
		resultPool.Put(result)
	}()
	n.search(x, y, &result)
	return result
}

func (n *Node) search(x, y float64, result *[]string) {
	if !n.Leaf {
		minX, maxX := x-n.Tree.radius, x+n.Tree.radius
		minY, maxY := y-n.Tree.radius, y+n.Tree.radius

		for _, son := range n.Child {
			if son.intersects(minX, minY) || son.intersects(maxX, minY) ||
				son.intersects(minX, maxY) || son.intersects(maxX, maxY) {
				son.search(x, y, result)
			}
		}
		return
	}

	n.Entities.Range(func(key, value interface{}) bool {
		*result = append(*result, value.(*Entity).Key)
		return true
	})
	return
}