add rebalance and fix range issues

threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt

@@ -150,6 +150,7 @@ class Document : IDocument {
     override fun getNodeByNodeId(nodeId: Long): INode<*>? = this.allNodes[nodeId]
 
     override fun findAt(vec: Vector3dc) = kdTree.find(vec)
+    override fun rebalance() = kdTree.rebalance()
 }
 
 

threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt

@@ -13,6 +13,7 @@ interface IDocument : EventDispatcher, INodeContainer {
     fun getNodeByNodeId(nodeId: Long): INode<*>?
     fun findAt(x: Number, y: Number, z: Number) = findAt(Vector3d(x, y, z))
     fun findAt(vec: Vector3dc): Collection<INode<*>>
+    fun rebalance()
 }
 
 inline fun <reified T : INode<T>> IDocument.createNode() = createNode(T::class)

threedom/src/main/kotlin/me/eater/threedom/utils/KDTree.kt

@@ -7,7 +7,7 @@ import me.eater.threedom.utils.joml.compareTo
 import me.eater.threedom.utils.joml.getTranslation
 import org.joml.Vector3dc
 
-class KDTree(private val document: IDocument, private val root: Node = Node(Vector3d(0, 0, 0))) {
+class KDTree(private val document: IDocument, private var root: Node = Node(Vector3d(0, 0, 0))) {
     private val nodeLocMap = mutableMapOf<Long, Vector3dc>()
 
     data class Node(
@@ -89,10 +89,10 @@ class KDTree(private val document: IDocument, private val root: Node = Node(Vect
                 val sorted = nodes.keys.sortedBy { it[axis] }
                 val median = sorted.size / 2
                 val selected = sorted[median]
-                val left = sorted.slice(0..median).toSet().takeIf { it.isNotEmpty() }?.let {
+                val left = sorted.slice(0 until median).toSet().takeIf { it.isNotEmpty() }?.let {
                     nodes.filterKeys(it::contains)
                 }?.let { create(it, depth + 1) }
-                val right = sorted.slice(median + 1..sorted.size).toSet().takeIf { it.isNotEmpty() }?.let {
+                val right = sorted.slice(median + 1 until sorted.size).toSet().takeIf { it.isNotEmpty() }?.let {
                     nodes.filterKeys(it::contains)
                 }?.let { create(it, depth + 1) }
                 return Node(selected, nodes[selected]?.toMutableSet() ?: mutableSetOf(), depth, left, right)
@@ -119,4 +119,8 @@ class KDTree(private val document: IDocument, private val root: Node = Node(Vect
         nodeLocMap[node.nodeId]?.let { root.remove(it, node.nodeId) }
         add(node)
     }
+
+    fun rebalance() {
+        root = Node.create(nodeLocMap.entries.groupBy({ it.value }) { it.key })
+    }
 }

threedom/src/test/kotlin/me/eater/test/threedom/dom/PositionTest.kt

@@ -38,6 +38,10 @@ class PositionTest : StringSpec({
         nodeTwo.model { setTranslation(-10, 20, 0) }
         nodeThree.model { setTranslation(0, 20, 0) }
 
+        doc.findAt(10, 0, 10) shouldHaveSingleElement node
+        doc.findAt(0, 20, 10) shouldHaveSingleElement nodeTwo
+        doc.findAt(0, 40, 10) shouldHaveSingleElement nodeThree
+        doc.rebalance()
         doc.findAt(10, 0, 10) shouldHaveSingleElement node
         doc.findAt(0, 20, 10) shouldHaveSingleElement nodeTwo
         doc.findAt(0, 40, 10) shouldHaveSingleElement nodeThree