go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/quad/tree.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package quad

// New returns a new QuadTree with a given set of options.
//
// Both fields are required,
func New[A any](opts ...Option) *Tree[A] {
	options := Options{
		MaxValuesPerQuad: 255,
		Center: Point{
			X: 1024,
			Y: 1024,
		},
		HalfDimensions: Dimension{
			Height: 1024,
			Width:  1024,
		},
	}
	for _, opt := range opts {
		opt(&options)
	}
	return &Tree[A]{
		maxValuesPerQuad:        options.MaxValuesPerQuad,
		values:                  make([]PointValue[A], 0, options.MaxValuesPerQuad),
		preallocateValueStorage: options.PreallocateValueStorage,
		Range: Range{
			Center:         options.Center,
			HalfDimensions: options.HalfDimensions,
		},
	}
}

// Tree creates a quad tree structure to organize
// values with given points for fast lookups.
type Tree[A any] struct {
	Range

	root                    *Tree[A]
	maxValuesPerQuad        int
	preallocateValueStorage bool

	depth  int
	zone   Zone
	values []PointValue[A]

	parent *Tree[A]
	nw     *Tree[A]
	ne     *Tree[A]
	sw     *Tree[A]
	se     *Tree[A]
}

// Option is a function that mutates QuadTreeOptions
type Option func(*Options)

// OptMaxValuesPerQuad sets the `MaxValuesPerQuad` field on the options.
func OptMaxValuesPerQuad(maxValuesPerQuad int) Option {
	return func(opts *Options) {
		opts.MaxValuesPerQuad = maxValuesPerQuad
	}
}

// OptCenter sets the `Center` field on the options.
func OptCenter(center Point) Option {
	return func(opts *Options) {
		opts.Center = center
	}
}

// OptHalfDimensions sets the `HalfDimensions` field on the options.
func OptHalfDimensions(halfDimensions Dimension) Option {
	return func(opts *Options) {
		opts.HalfDimensions = halfDimensions
	}
}

// OptPreallocateValueStorage sets the `PreallocateValueStorage` field on the options.
func OptPreallocateValueStorage(preallocateValueStorage bool) Option {
	return func(opts *Options) {
		opts.PreallocateValueStorage = preallocateValueStorage
	}
}

// Options is options for a quad tree's constructor.
type Options struct {
	MaxValuesPerQuad        int
	Center                  Point
	HalfDimensions          Dimension
	PreallocateValueStorage bool
}

// Insert adds a new point with a given value.
func (qt *Tree[A]) Insert(p Point, v A) bool {
	if !qt.ContainsPoint(p) {
		return false
	}
	if len(qt.values) <= qt.maxValuesPerQuad && qt.nw == nil {
		qt.values = append(qt.values, PointValue[A]{
			Point: p,
			Zone:  qt.zone,
			Value: v,
		})
		return true
	}
	if qt.nw == nil {
		qt.subdivide(qt.preallocateValueStorage)
	}
	if qt.nw.Insert(p, v) {
		return true
	} else if qt.ne.Insert(p, v) {
		return true
	} else if qt.sw.Insert(p, v) {
		return true
	} else if qt.se.Insert(p, v) {
		return true
	}
	return false
}

// QueryRange queries the tree for all points within a given
// range as denoted by a given QuadRange.
func (qt *Tree[A]) QueryZone(zone Zone) []PointValue[A] {
	if zone == "" {
		return qt.values
	}

	decoded := DecodeZoneHigh(zone)
	switch decoded {
	case SW:
		if qt.sw == nil {
			return nil
		}
		return qt.sw.QueryZone(ShiftZone(zone))
	case SE:
		if qt.se == nil {
			return nil
		}
		return qt.se.QueryZone(ShiftZone(zone))
	case NW:
		if qt.nw == nil {
			return nil
		}
		return qt.nw.QueryZone(ShiftZone(zone))
	case NE:
		if qt.ne == nil {
			return nil
		}
		return qt.ne.QueryZone(ShiftZone(zone))
	}
	return nil
}

// QueryRange queries the tree for all points within a given
// range as denoted by a given QuadRange.
func (qt *Tree[A]) QueryRange(qr Range) []PointValue[A] {
	if !qt.Intersects(qr) {
		return nil
	}
	var output []PointValue[A]
	for _, p := range qt.values {
		if qr.ContainsPoint(p.Point) {
			output = append(output, p)
		}
	}
	if qt.nw == nil {
		return output
	}
	output = append(output, qt.nw.QueryRange(qr)...)
	output = append(output, qt.ne.QueryRange(qr)...)
	output = append(output, qt.sw.QueryRange(qr)...)
	output = append(output, qt.se.QueryRange(qr)...)
	return output
}

// QueryPoint gets all the values in the quad tree boundary
// as found by a single point within the smallest boundary of the tree.
func (qt *Tree[A]) QueryPoint(p Point) []PointValue[A] {
	if !qt.ContainsPoint(p) {
		return nil
	}
	if qt.sw == nil {
		return qt.values
	}
	if qt.sw.ContainsPoint(p) {
		return qt.sw.QueryPoint(p)
	}
	if qt.se.ContainsPoint(p) {
		return qt.se.QueryPoint(p)
	}
	if qt.nw.ContainsPoint(p) {
		return qt.nw.QueryPoint(p)
	}
	if qt.ne.ContainsPoint(p) {
		return qt.ne.QueryPoint(p)
	}
	return nil
}

// ZoneForPoint gets the zone identifier for a given point.
func (qt *Tree[A]) ZoneForPoint(p Point) (zone Zone, ok bool) {
	if !qt.ContainsPoint(p) {
		return
	}
	if qt.sw == nil {
		zone = qt.zone
		ok = true
		return
	}
	if qt.sw.ContainsPoint(p) {
		return qt.sw.ZoneForPoint(p)
	}
	if qt.se.ContainsPoint(p) {
		return qt.se.ZoneForPoint(p)
	}
	if qt.nw.ContainsPoint(p) {
		return qt.nw.ZoneForPoint(p)
	}
	if qt.ne.ContainsPoint(p) {
		return qt.ne.ZoneForPoint(p)
	}
	return
}

// Values yields the values in the current tree node.
func (qt *Tree[A]) Values() (output []PointValue[A]) {
	output = make([]PointValue[A], len(qt.values))
	copy(output, qt.values)
	return
}

// ValuesAll yields all the values in the tree.
func (qt *Tree[A]) ValuesAll() (output []PointValue[A]) {
	qt.VisitDepth(func(n *Tree[A]) {
		output = append(output, n.values...)
	})
	return
}

// VisitDepth visits all the nodes of a quad tree
// in depth first order according to the z-traversal rules.
//
// E.g we visit SW, SE, NW, NE for each node.
func (qt *Tree[A]) VisitDepth(fn func(*Tree[A])) {
	fn(qt)
	if qt.se != nil {
		qt.sw.VisitDepth(fn)
		qt.se.VisitDepth(fn)
		qt.nw.VisitDepth(fn)
		qt.ne.VisitDepth(fn)
	}
}

// subdivide creates 4 separate sub-trees and moves
// values into those trees according to their points
// and the sub-trees new bounds.
func (qt *Tree[A]) subdivide(preallocate bool) {
	newDimensions := Dimension{
		Width:  qt.HalfDimensions.Width / 2,
		Height: qt.HalfDimensions.Height / 2,
	}
	var root *Tree[A]
	if qt.root != nil {
		root = qt.root
	} else {
		root = qt
	}

	var swValues []PointValue[A]
	if preallocate {
		swValues = make([]PointValue[A], 0, qt.maxValuesPerQuad)
	}
	qt.sw = &Tree[A]{
		root:             root,
		parent:           qt,
		maxValuesPerQuad: qt.maxValuesPerQuad,
		Range: Range{
			Center: Point{
				X: qt.Center.X - newDimensions.Width,
				Y: qt.Center.Y + newDimensions.Height,
			},
			HalfDimensions: newDimensions,
		},
		values: swValues,
		depth:  qt.depth + 1,
		zone:   AppendZone(qt.zone, SW),
	}

	var seValues []PointValue[A]
	if preallocate {
		seValues = make([]PointValue[A], 0, qt.maxValuesPerQuad)
	}
	qt.se = &Tree[A]{
		root:             root,
		parent:           qt,
		maxValuesPerQuad: qt.maxValuesPerQuad,
		Range: Range{
			Center: Point{
				X: qt.Center.X + newDimensions.Width,
				Y: qt.Center.Y + newDimensions.Height,
			},
			HalfDimensions: newDimensions,
		},
		values: seValues,
		depth:  qt.depth + 1,
		zone:   AppendZone(qt.zone, SE),
	}

	var nwValues []PointValue[A]
	if preallocate {
		nwValues = make([]PointValue[A], 0, qt.maxValuesPerQuad)
	}
	qt.nw = &Tree[A]{
		root:             root,
		parent:           qt,
		maxValuesPerQuad: qt.maxValuesPerQuad,
		Range: Range{
			Center: Point{
				X: qt.Center.X - newDimensions.Width,
				Y: qt.Center.Y - newDimensions.Height,
			},
			HalfDimensions: newDimensions,
		},
		values: nwValues,
		depth:  qt.depth + 1,
		zone:   AppendZone(qt.zone, NW),
	}

	var neValues []PointValue[A]
	if preallocate {
		neValues = make([]PointValue[A], 0, qt.maxValuesPerQuad)
	}
	qt.ne = &Tree[A]{
		root:             root,
		parent:           qt,
		maxValuesPerQuad: qt.maxValuesPerQuad,
		Range: Range{
			Center: Point{
				X: qt.Center.X + newDimensions.Width,
				Y: qt.Center.Y - newDimensions.Height,
			},
			HalfDimensions: newDimensions,
		},
		values: neValues,
		depth:  qt.depth + 1,
		zone:   AppendZone(qt.zone, NE),
	}

	// redistribute the values
	for _, pv := range qt.values {
		if qt.sw.Insert(pv.Point, pv.Value) {
			continue
		} else if qt.se.Insert(pv.Point, pv.Value) {
			continue
		} else if qt.nw.Insert(pv.Point, pv.Value) {
			continue
		} else if qt.ne.Insert(pv.Point, pv.Value) {
			continue
		}
	}
	qt.values = nil
}