My first attempt at implementing John Ratcliffe's Sphere Trees in go

spheretree.go

package main

import (
	"fmt"
	"github.com/faiface/pixel"
	"github.com/faiface/pixel/imdraw"
	"github.com/faiface/pixel/pixelgl"
	"github.com/faiface/pixel/text"
	"golang.org/x/image/colornames"
	"golang.org/x/image/font/basicfont"
	"image/color"
	"math"
	"math/rand"
	"time"
)

var ZeroVector3 = Vector3{}

// Zero is a utility function that returns the zero value for any type
func Zero[T any]() T {
	return *new(T)
}

// FreeList is a structure that holds any kind of object.
// When an entry is erased, it marks that slot as free so that new entries can be filled in to
// the existing allocated memory.
type FreeList[T any] struct {
	data      []freeListEntry[T]
	FirstFree int
}

type freeListEntry[T any] struct {
	element  T
	nextFree int
}

func NewFreeList[T any]() FreeList[T] {
	return FreeList[T]{
		FirstFree: -1,
	}
}

func (f *FreeList[T]) Insert(element T) int {
	// if there is a free entry
	if f.FirstFree != -1 {
		// get the index of the free entry
		index := f.FirstFree
		// set the next free entry from this index
		f.FirstFree = f.data[index].nextFree
		f.data[index].nextFree = 0
		// set the data at current free index
		f.data[index].element = element
		return index
	}
	// insert at the end of the list
	f.data = append(f.data, freeListEntry[T]{
		element:  element,
		nextFree: 0,
	})
	return len(f.data) - 1
}

func (f *FreeList[T]) Set(n int, element T) {
	f.data[n].element = element
}

func (f *FreeList[T]) Erase(n int) {
	// set the old next free index to this node
	f.data[n].nextFree = f.FirstFree
	// set the first free index to this nodes next free Raw
	f.FirstFree = n
}

func (f *FreeList[T]) Clear() {
	f.data = f.data[:0]
	f.FirstFree = -1
}

func (f *FreeList[T]) Get(n int) T {
	return f.data[n].element
}

func (f *FreeList[T]) Len() int {
	return len(f.data)
}

// Queue is your standard First In First Out queue.
type Queue[T any] struct {
	data []T
}

func (q *Queue[T]) Push(e T) {
	q.data = append(q.data, e)
}

func (q *Queue[T]) Pop() (T, bool) {
	if len(q.data) == 0 {
		return Zero[T](), false
	}
	e := q.data[0]
	q.data = q.data[1:]
	return e, true
}

func (q *Queue[T]) Peek(i int) T {
	return q.data[i]
}

func (q *Queue[T]) Len() int {
	return len(q.data)
}

// Vector3 represents a point or a line in 3d space
type Vector3 struct {
	X float64
	Y float64
	Z float64
}

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

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

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

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

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

func (v Vector3) Distance(v2 Vector3) float64 {
	return math.Sqrt((v2.X-v.X)*(v2.X-v.X) + (v2.Y-v.Y)*(v2.Y-v.Y) + (v2.Z-v.Z)*(v2.Z-v.Z))
}

// Sphere is literally a sphere in 3d space
type Sphere struct {
	center Vector3
	radius float64
}

func NewSphere(center Vector3, radius float64) Sphere {
	return Sphere{
		center: center,
		radius: radius,
	}
}

func (s Sphere) IntersectsSphere(s2 Sphere) bool {
	return s.center.Distance(s2.center) < s.radius+s2.radius
}

func (s Sphere) ContainsSphere(s2 Sphere) bool {
	return s.radius >= s.center.Distance(s2.center)+s2.radius
}

// SphereTree is a space partitioning data structure that contains spheres inside larger super sphere
type SphereTree struct {
	spheres       FreeList[SphereEntry]
	integrateFIFO Queue[int]
	recomputeFIFO Queue[int]
	maxSize       float64
	gravy         float64
}

type SphereEntry struct {
	ID         uint64
	sphere     Sphere
	parent     int
	firstChild int
	next       int
}

func NewSphereTree(rootSphere Sphere, maxSize, gravy float64) *SphereTree {
	st := &SphereTree{
		spheres: FreeList[SphereEntry]{FirstFree: -1},
		maxSize: maxSize,
		gravy:   gravy,
	}

	st.spheres.Insert(SphereEntry{sphere: rootSphere, firstChild: -1, next: -1})

	return st
}

func (st *SphereTree) Insert(id uint64, sphere Sphere) int {
	entry := SphereEntry{ID: id, sphere: sphere, firstChild: -2, next: -1}
	// insert entry into list
	i := st.spheres.Insert(entry)
	// set the entry to be integrated into the tree, don't do it now
	st.QueueIntegrate(i)
	return i
}

func (st *SphereTree) Remove(entryID int) {
	entry := st.spheres.Get(entryID)
	st.spheres.Erase(entryID)

	// remove this entry from the parent's linked list
	st.removeChild(entry.parent, entryID)
}

func (st *SphereTree) Move(entryID int, sphere Sphere) {
	entry := st.spheres.Get(entryID)
	entry.sphere = sphere
	st.spheres.Set(entryID, entry)

	parentID := entry.parent
	parent := st.spheres.Get(parentID)
	// parent spheres still contains
	if parent.sphere.ContainsSphere(entry.sphere) {
		//st.QueueRecompute(parentID) // TODO: maybe dont
		return
	}

	// entry has broken out of sphere
	// so detach it and queue it for integration
	// and recompute its old parent
	st.removeChild(parentID, entryID)
	st.QueueIntegrate(entryID)
}

func (st *SphereTree) Integrate() {
	for {
		integrateCandidateID, ok := st.integrateFIFO.Pop()
		if !ok {
			break
		}
		st.integrate(integrateCandidateID)
	}
}

func (st *SphereTree) integrate(entryID int) {
	integrateCandidate := st.spheres.Get(entryID)

	containsUs := -1
	nearest := -1
	nearestDist := math.MaxFloat64

	// look through all super spheres for candidate parent
	// look for super spheres that fully contain the candidate or find the closet one
	rootSphere := st.spheres.Get(0) // root
	superSphereIndex := rootSphere.firstChild
	for {
		if superSphereIndex < 0 {
			break
		}
		superSphere := st.spheres.Get(superSphereIndex)
		if superSphere.sphere.ContainsSphere(integrateCandidate.sphere) {
			containsUs = superSphereIndex
			break
		}

		// TODO: Verify this
		dist := superSphere.sphere.center.Distance(integrateCandidate.sphere.center) + integrateCandidate.sphere.radius - superSphere.sphere.radius
		if dist < nearestDist {
			nearest = superSphereIndex
			nearestDist = dist
		}

		superSphereIndex = superSphere.next
	}

	// if a super sphere contains it, just insert it!
	if containsUs != -1 {
		st.addChild(containsUs, entryID)
		return
	}

	// check to see if the nearest sphere can grow to contain us
	if nearest != -1 {
		parent := st.spheres.Get(nearest)
		newSize := nearestDist + parent.sphere.radius
		if newSize <= st.maxSize {
			parent.sphere.radius = newSize + st.gravy
			st.spheres.Set(nearest, parent)
			st.addChild(nearest, entryID)
			return
		}
	}

	// we'll have to make a new super sphere
	parent := SphereEntry{
		sphere:     integrateCandidate.sphere,
		firstChild: -1,
		parent:     0,
	}
	parent.sphere.radius += st.gravy
	rootSphere = st.spheres.Get(0) // root
	parent.next = rootSphere.firstChild
	parentID := st.spheres.Insert(parent)
	rootSphere.firstChild = parentID
	st.spheres.Set(0, rootSphere)

	st.addChild(parentID, entryID)
}

func (st *SphereTree) Recompute() {
	for {
		recomputeCandidateID, ok := st.recomputeFIFO.Pop()
		if !ok {
			break
		}
		st.recompute(recomputeCandidateID)
	}
}

func (st *SphereTree) recompute(superSphereID int) {
	superSphere := st.spheres.Get(superSphereID)

	if superSphere.firstChild == -1 {
		st.removeChild(0, superSphereID)
		st.spheres.Erase(superSphereID)
		return
	}

	var childCount int
	var total Vector3
	childIndex := superSphere.firstChild
	for {
		if childIndex == -1 {
			break
		}

		childCount++
		child := st.spheres.Get(childIndex)
		total = total.Add(child.sphere.center)
		childIndex = child.next
	}
	recip := 1.0 / float64(childCount)
	oldCenter := superSphere.sphere.center
	superSphere.sphere.center = total.Multiply(recip)

	newRadius := 0.0
	childIndex = superSphere.firstChild
	for {
		if childIndex == -1 {
			break
		}
		child := st.spheres.Get(childIndex)
		radius := superSphere.sphere.center.Distance(child.sphere.center) + child.sphere.radius
		if radius > newRadius {
			newRadius = radius
			if newRadius+st.gravy > superSphere.sphere.radius {
				superSphere.sphere.center = oldCenter
				return
			}
		}
		childIndex = child.next
	}
	superSphere.sphere.radius = newRadius + st.gravy
	st.spheres.Set(superSphereID, superSphere)

	// RINSE: look for circle that can own this one
	root := st.spheres.Get(0)
	possibleParentID := root.firstChild
	for {
		if possibleParentID == 0 {
			panic("no")
		}
		if possibleParentID == -1 {
			break
		}

		if possibleParentID == superSphereID {
			possibleParentID = superSphere.next
			continue
		}

		possibleParent := st.spheres.Get(possibleParentID)
		if possibleParent.sphere.ContainsSphere(superSphere.sphere) {
			// find last child of possible parent

			childIndex := superSphere.firstChild
			for {
				if childIndex == -1 {
					break
				}

				child := st.spheres.Get(childIndex)
				next := child.next
				child.next = possibleParent.firstChild
				child.parent = possibleParentID
				possibleParent.firstChild = childIndex
				st.spheres.Set(childIndex, child)

				childIndex = next
			}

			superSphere.firstChild = -1
			st.spheres.Set(superSphereID, superSphere)
			st.spheres.Set(possibleParentID, possibleParent)
			st.QueueRecompute(superSphereID)
			st.QueueRecompute(possibleParentID)

			break
		}

		possibleParentID = possibleParent.next
	}
}

func (st *SphereTree) Walk(f func(s Sphere, isSuper bool)) {
	root := st.spheres.Get(0)

	childIndex := root.firstChild
	for {
		if childIndex == -1 {
			break
		}
		child := st.spheres.Get(childIndex)
		entryIndex := child.firstChild
		for {
			if entryIndex == -1 {
				break
			}
			entry := st.spheres.Get(entryIndex)
			f(entry.sphere, false)
			entryIndex = entry.next
		}
		childIndex = child.next
	}

	childIndex = root.firstChild
	for {
		if childIndex == -1 {
			break
		}
		child := st.spheres.Get(childIndex)
		f(child.sphere, true)
		childIndex = child.next
	}
}

func (st *SphereTree) QueueIntegrate(entryID int) {
	for i := 0; i < st.integrateFIFO.Len(); i++ {
		if st.integrateFIFO.Peek(i) == entryID {
			return
		}
	}
	st.integrateFIFO.Push(entryID)
}

func (st *SphereTree) QueueRecompute(parentID int) {
	if parentID == 0 {
		return
	}
	for i := 0; i < st.recomputeFIFO.Len(); i++ {
		if st.recomputeFIFO.Peek(i) == parentID {
			return
		}
	}
	st.recomputeFIFO.Push(parentID)
}

func (st *SphereTree) GetEntries(selectionSphere Sphere) []SphereEntry {
	var results []SphereEntry

	root := st.spheres.Get(0)
	sphereIndex := root.firstChild

	for {
		if sphereIndex == -1 {
			break
		}
		sphere := st.spheres.Get(sphereIndex)
		if selectionSphere.IntersectsSphere(sphere.sphere) {
			childIndex := sphere.firstChild
			for {
				if childIndex == -1 {
					break
				}

				child := st.spheres.Get(childIndex)

				if selectionSphere.IntersectsSphere(child.sphere) {
					results = append(results, child)
				}

				childIndex = child.next
			}
		}

		sphereIndex = sphere.next
	}

	return results
}

func (st *SphereTree) addChild(parentID, sphereID int) {
	parent := st.spheres.Get(parentID)
	entry := st.spheres.Get(sphereID)

	entry.parent = parentID
	entry.next = parent.firstChild
	parent.firstChild = sphereID

	st.spheres.Set(parentID, parent)
	st.spheres.Set(sphereID, entry)

	st.QueueRecompute(parentID)
}

func (st *SphereTree) removeChild(parentID, entryID int) {
	parent := st.spheres.Get(parentID)
	entry := st.spheres.Get(entryID)

	childIndex := parent.firstChild
	if childIndex == entryID {
		parent.firstChild = entry.next
	} else {
		for {
			if childIndex == -1 {
				break
			}
			child := st.spheres.Get(childIndex)
			if child.next == entryID {
				child.next = entry.next
				st.spheres.Set(childIndex, child)
				break
			}
			childIndex = child.next
		}
	}

	st.spheres.Set(parentID, parent)
	st.spheres.Set(entryID, entry)

	st.QueueRecompute(parentID)
}

type actor struct {
	id        uint64
	sphere    Sphere
	entryID   int
	attractor Vector3
	speed     float64
	color     color.Color
}

func main() {
	pixelgl.Run(run)
}

func run() {
	rand.Seed(time.Now().UnixMilli())

	cfg := pixelgl.WindowConfig{
		Title:     "Pixel Rocks!",
		Bounds:    pixel.R(0, 0, 1920, 1080),
		VSync:     true,
		Resizable: true,
		Maximized: false,
	}

	win, err := pixelgl.NewWindow(cfg)
	if err != nil {
		panic(err)
	}

	atlas := text.NewAtlas(basicfont.Face7x13, text.ASCII)
	moustTxt := text.New(pixel.V(50, win.Bounds().H()-50), atlas)

	var actors []actor
	var attractors []Vector3
	for i := 0; i < 19; i++ {
		attractors = append(attractors, Vector3{X: rand.Float64() * win.Bounds().W(), Y: rand.Float64() * win.Bounds().H()})
	}

	st := NewSphereTree(Sphere{center: Vector3{X: win.Bounds().W(), Y: win.Bounds().H() / 2.0}, radius: 500}, 100, 10)

	imd := imdraw.New(nil)
	for !win.Closed() {
		// updates
		if len(actors) < 4000 {
			for i := 0; i < 100; i++ {
				sphere := Sphere{center: Vector3{X: rand.Float64() * win.Bounds().W(), Y: rand.Float64() * win.Bounds().H()}, radius: 1.0}
				entryID := st.Insert(uint64(len(actors)), sphere)

				var attractor Vector3
				c := colornames.Deepskyblue
				if len(actors) > 2000 {
					c = colornames.Red
					if i%3 == 0 {
						c = colornames.Purple
					}
					if i%3 == 1 {
						c = colornames.Hotpink
					}
					attractor = attractors[rand.Intn(len(attractors))]
				}

				actors = append(actors, actor{
					id:        uint64(len(actors)),
					sphere:    sphere,
					entryID:   entryID,
					attractor: attractor,
					speed:     1,
					color:     c,
				})
			}
		} else {
			for i := 0; i < len(actors); i++ {
				if actors[i].attractor == ZeroVector3 {
					continue
				}
				delta := actors[i].attractor.Sub(actors[i].sphere.center)
				if delta.Magnitude() < 1.0 {
					actors[i].sphere.center = Vector3{X: rand.Float64() * win.Bounds().W(), Y: rand.Float64() * win.Bounds().H()}
					delta = actors[i].attractor.Sub(actors[i].sphere.center)
				}
				actors[i].sphere.center = actors[i].sphere.center.Add(delta.Normalize().Multiply(actors[i].speed))
				st.Move(actors[i].entryID, actors[i].sphere)
			}
		}

		st.Integrate()
		st.Recompute()

		// renders

		imd.Clear()

		for i := range actors {
			imd.Color = actors[i].color
			imd.Push(pixel.V(actors[i].sphere.center.X, actors[i].sphere.center.Y))
			imd.Circle(actors[i].sphere.radius, 1)
		}

		imd.Color = colornames.Navy
		st.Walk(func(s Sphere, isSuper bool) {
			if !isSuper {
				return
			}
			imd.Push(pixel.V(s.center.X, s.center.Y))
			imd.Circle(s.radius, 1)
		})
		for _, a := range attractors {
			imd.Color = colornames.Limegreen
			imd.Push(pixel.V(a.X, a.Y))
			imd.Circle(4, 2)
		}

		imd.Color = colornames.Yellowgreen
		//now := time.Now()
		entries := st.GetEntries(Sphere{
			center: Vector3{X: win.MousePosition().X, Y: win.MousePosition().Y},
			radius: 40,
		})
		//fmt.Println(time.Since(now))
		for i := range entries {
			imd.Push(pixel.V(entries[i].sphere.center.X, entries[i].sphere.center.Y))
			imd.Circle(entries[i].sphere.radius, 1)
		}
		imd.Push(win.MousePosition())
		imd.Circle(40, 1)

		moustTxt.Clear()
		fmt.Fprintf(moustTxt, "Selected: %d", len(entries))

		win.Clear(color.Black)
		imd.Draw(win)
		moustTxt.Draw(win, pixel.IM)
		win.Update()
	}
}