Add region search

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

@@ -151,6 +151,7 @@ class Document : IDocument {
 
     override fun findAt(vec: Vector3dc) = kdTree.find(vec)
     override fun rebalance() = kdTree.rebalance()
+    override fun findInRegion(pointA: Vector3dc, pointB: Vector3dc) = kdTree.findInRegion(pointA, pointB)
 }
 
 

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

@@ -14,6 +14,7 @@ interface IDocument : EventDispatcher, INodeContainer {
     fun findAt(x: Number, y: Number, z: Number) = findAt(Vector3d(x, y, z))
     fun findAt(vec: Vector3dc): Collection<INode<*>>
     fun rebalance()
+    fun findInRegion(pointA: Vector3dc, pointB: Vector3dc): Sequence<INode<*>>
 }
 
 inline fun <reified T : INode<T>> IDocument.createNode() = createNode(T::class)

threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeClassListUpdate.kt

@@ -1,10 +1,14 @@
 package me.eater.threedom.dom.event
 
 import me.eater.threedom.dom.INode
+import me.eater.threedom.kapt.EventName
 
 sealed class NodeClassListUpdate {
     abstract val node: INode<*>
 
-    class Removed(val classNames: Set<String>, override val node: INode<*>): NodeClassListUpdate()
-    class Added(val classNames: Set<String>, override val node: INode<*>): NodeClassListUpdate()
+    @EventName("NodeClassesRemoved")
+    class Removed(val classNames: Set<String>, override val node: INode<*>) : NodeClassListUpdate()
+
+    @EventName("NodeClassesAdded")
+    class Added(val classNames: Set<String>, override val node: INode<*>) : NodeClassListUpdate()
 }

threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeIDUpdate.kt

@@ -1,5 +1,7 @@
 package me.eater.threedom.dom.event
 
 import me.eater.threedom.dom.INode
+import me.eater.threedom.kapt.EventName
 
+@EventName("NodeIDUpdate")
 data class NodeIDUpdate(val node: INode<*>, val old: String?, val new: String?)

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

@@ -10,16 +10,27 @@ import org.joml.Vector3dc
 class KDTree(private val document: IDocument, private var root: Node = Node(Vector3d(0, 0, 0))) {
     private val nodeLocMap = mutableMapOf<Long, Vector3dc>()
 
-    data class Node(
+    class Node(
         val vertex: Vector3dc = Vector3d(0, 0, 0),
         val nodeIds: MutableSet<Long> = mutableSetOf(),
         val depth: Long = 0,
-        var left: Node? = null,
-        var right: Node? = null
+        val branches: Array<Node?> = Array(3) { null }
     ) {
         val axis: Int
             get() = (depth % 3).toInt()
 
+        val median: Double
+            get() = vertex[axis]
+
+        val left: Node?
+            get() = branches[0]
+
+        val middle: Node?
+            get() = branches[1]
+
+        val right: Node?
+            get() = branches[2]
+
         fun add(translation: Vector3dc, nodeId: Long) {
             var current = this
 
@@ -29,20 +40,12 @@ class KDTree(private val document: IDocument, private var root: Node = Node(Vect
                     return
                 }
 
-                if (translation[current.axis] < current.vertex[current.axis] || (translation[current.axis] == current.vertex[current.axis] && translation < current.vertex)) {
-                    if (current.left == null) {
-                        current.left = Node(translation, mutableSetOf(nodeId), current.depth + 1)
-                        return
-                    } else {
-                        current = current.left!!
-                    }
+                val branch = 1 + translation[current.axis].compareTo(current.vertex[current.axis])
+                if (current.branches[branch] == null) {
+                    current.branches[branch] = Node(translation, mutableSetOf(nodeId), current.depth + 1)
+                    return
                 } else {
-                    if (current.right == null) {
-                        current.right = Node(translation, mutableSetOf(nodeId), current.depth + 1)
-                        return
-                    } else {
-                        current = current.right!!
-                    }
+                    current = current.branches[branch]!!
                 }
             }
         }
@@ -51,6 +54,30 @@ class KDTree(private val document: IDocument, private var root: Node = Node(Vect
             find(translation)?.nodeIds?.remove(nodeId)
         }
 
+        fun findInRegion(pointA: Vector3dc, pointB: Vector3dc) = sequence<Long> {
+            var current: Node? = this@Node
+            val selection = mutableListOf<Node>()
+            while (current != null) {
+                if (pointA <= current.vertex && current.vertex <= pointB) {
+                    yieldAll(current.nodeIds)
+                }
+
+                if (pointA[current.axis] < current.median) {
+                    current.left?.let(selection::add)
+                }
+
+                if (pointA[current.axis] <= current.median && current.median <= pointB[current.axis]) {
+                    current.middle?.let(selection::add)
+                }
+
+                if (pointB[current.axis] > current.median) {
+                    current.right?.let(selection::add)
+                }
+
+                current = selection.firstOrNull()?.apply { selection.removeAt(0) }
+            }
+        }
+
         fun find(translation: Vector3dc): Node? {
             var current: Node? = this
             while (current != null) {
@@ -58,12 +85,7 @@ class KDTree(private val document: IDocument, private var root: Node = Node(Vect
                     return current
                 }
 
-                current =
-                    if (translation[current.axis] < current.vertex[current.axis] || (translation[current.axis] == current.vertex[current.axis] && translation < current.vertex)) {
-                        current.left
-                    } else {
-                        current.right
-                    }
+                current = current.branches[1 + translation[current.axis].compareTo(current.vertex[current.axis])]
             }
 
             return null
@@ -86,16 +108,22 @@ class KDTree(private val document: IDocument, private var root: Node = Node(Vect
                 }
 
                 val axis: Int = (depth % 3).toInt()
-                val sorted = nodes.keys.sortedBy { it[axis] }
+                val sorted = nodes.keys.sortedBy { it[axis] }.toMutableList()
                 val median = sorted.size / 2
                 val selected = sorted[median]
-                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 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)
+
+                val branches: Array<MutableMap<Vector3dc, Collection<Long>>> =
+                    arrayOf(mutableMapOf(), mutableMapOf(), mutableMapOf())
+                for (item in sorted) {
+                    nodes[item]?.let { branches[1 + item[axis].compareTo(selected[axis])][item] = it }
+                }
+
+                return Node(
+                    selected,
+                    nodes[selected]?.toMutableSet() ?: mutableSetOf(),
+                    depth,
+                    branches.map { Node.create(it, depth + 1) }.toTypedArray()
+                )
             }
         }
     }
@@ -115,6 +143,9 @@ class KDTree(private val document: IDocument, private var root: Node = Node(Vect
 
     fun find(vertex: Vector3dc) = root.find(vertex)?.nodeIds?.mapNotNull(document::getNodeByNodeId) ?: emptyList()
 
+    fun findInRegion(pointA: Vector3dc, pointB: Vector3dc) =
+        root.findInRegion(pointA, pointB).mapNotNull(document::getNodeByNodeId)
+
     fun update(node: INode<*>) {
         nodeLocMap[node.nodeId]?.let { root.remove(it, node.nodeId) }
         add(node)

threedom/src/main/kotlin/me/eater/threedom/utils/joml/JOML.kt

@@ -18,7 +18,7 @@ operator fun Matrix4d.times(rhs: Matrix4dc) = mul(rhs)
 fun Vector3d(x: Number, y: Number, z: Number) = Vector3d(x.toDouble(), y.toDouble(), z.toDouble())
 
 operator fun Vector3dc.compareTo(rhs: Vector3dc): Int {
-    for (i in 0..3) {
+    for (i in 0 until 3) {
         val c = this[i].compareTo(rhs[i])
         if (c != 0) {
             return c

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

@@ -1,7 +1,7 @@
 package me.eater.test.threedom.dom
 
 import io.kotest.core.spec.style.StringSpec
-import io.kotest.matchers.collections.shouldHaveSingleElement
+import io.kotest.matchers.collections.*
 import io.kotest.matchers.shouldBe
 import me.eater.threedom.dom.Document
 import me.eater.threedom.dom.IDocument
@@ -49,4 +49,25 @@ class PositionTest : StringSpec({
         doc.findAt(-10, 20, 0) shouldHaveSingleElement nodeTwo
         doc.findAt(-10, 40, 0) shouldHaveSingleElement nodeThree
     }
+
+    "ensure in selection works" {
+        val doc: IDocument = Document()
+        val node = doc.createNode<PlainNode>()
+        val nodeTwo = doc.createNode<PlainNode>()
+        val nodeThree = doc.createNode<PlainNode>()
+        doc.addNode(node)
+        node.addNode(nodeTwo)
+        nodeTwo.addNode(nodeThree)
+        node.model { setTranslation(10, 0, 10) }
+        nodeTwo.model { setTranslation(-10, 20, 0) }
+        nodeThree.model { setTranslation(0, 20, 0) }
+
+        val result = doc.findInRegion(Vector3d(0, 0, 0), Vector3d(0, 20, 20)).toList()
+        result shouldHaveSize 1
+        result shouldHaveSingleElement nodeTwo
+
+        val resultTwo = doc.findInRegion(Vector3d(0, 0, 0), Vector3d(0, 40, 20)).toList()
+        resultTwo shouldHaveSize 2
+        resultTwo shouldContainExactly listOf(nodeTwo, nodeThree)
+    }
 })