From ff5c65765e1338740063a3cc77f1d88d3f548c85 Mon Sep 17 00:00:00 2001
From: eater <=@eater.me>
Date: Sun, 17 May 2020 16:30:04 +0200
Subject: [PATCH] add rebalance and fix range issues

---
 .../src/main/kotlin/me/eater/threedom/dom/Document.kt  |  1 +
 .../src/main/kotlin/me/eater/threedom/dom/IDocument.kt |  1 +
 .../src/main/kotlin/me/eater/threedom/utils/KDTree.kt  | 10 +++++++---
 .../kotlin/me/eater/test/threedom/dom/PositionTest.kt  |  4 ++++
 4 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt b/threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt
index 23bfbc2..8f54b12 100644
--- a/threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt
+++ b/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()
 }
 
 
diff --git a/threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt b/threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt
index 718f184..1f8c0b0 100644
--- a/threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt
+++ b/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)
diff --git a/threedom/src/main/kotlin/me/eater/threedom/utils/KDTree.kt b/threedom/src/main/kotlin/me/eater/threedom/utils/KDTree.kt
index 70ba003..4614870 100644
--- a/threedom/src/main/kotlin/me/eater/threedom/utils/KDTree.kt
+++ b/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 })
+    }
 }
diff --git a/threedom/src/test/kotlin/me/eater/test/threedom/dom/PositionTest.kt b/threedom/src/test/kotlin/me/eater/test/threedom/dom/PositionTest.kt
index dfeec54..aae2f11 100644
--- a/threedom/src/test/kotlin/me/eater/test/threedom/dom/PositionTest.kt
+++ b/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