add findInRange

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

@@ -152,6 +152,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)
+    override fun findInRange(origin: Vector3dc, range: Number) = kdTree.findInRange(origin, range)
 }
 
 

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

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

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

@@ -54,12 +54,20 @@ 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> {
+        fun findInRange(origin: Vector3dc, range: Number): Sequence<Node> {
+            val rangeD = range.toDouble()
+            val pointA = Vector3d(origin.x() - rangeD, origin.y() - rangeD, origin.z() - rangeD)
+            val pointB = Vector3d(origin.x() + rangeD, origin.y() + rangeD, origin.z() + rangeD)
+
+            return findInRegion(pointA, pointB).filter { it.vertex.distance(origin) <= rangeD }
+        }
+
+        fun findInRegion(pointA: Vector3dc, pointB: Vector3dc) = sequence<Node> {
             var current: Node? = this@Node
             val selection = mutableListOf<Node>()
             while (current != null) {
                 if (pointA <= current.vertex && current.vertex <= pointB) {
-                    yieldAll(current.nodeIds)
+                    yield(current)
                 }
 
                 if (pointA[current.axis] < current.median) {
@@ -143,8 +151,11 @@ 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 findInRange(origin: Vector3dc, range: Number) =
+        root.findInRange(origin, range).flatMap { it.nodeIds.asSequence() }.mapNotNull(document::getNodeByNodeId)
+
     fun findInRegion(pointA: Vector3dc, pointB: Vector3dc) =
-        root.findInRegion(pointA, pointB).mapNotNull(document::getNodeByNodeId)
+        root.findInRegion(pointA, pointB).flatMap { it.nodeIds.asSequence() }.mapNotNull(document::getNodeByNodeId)
 
     fun update(node: INode<*>) {
         nodeLocMap[node.nodeId]?.let { root.remove(it, node.nodeId) }

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

@@ -1,7 +1,9 @@
 package me.eater.test.threedom.dom
 
 import io.kotest.core.spec.style.StringSpec
-import io.kotest.matchers.collections.*
+import io.kotest.matchers.collections.shouldContainExactly
+import io.kotest.matchers.collections.shouldHaveSingleElement
+import io.kotest.matchers.collections.shouldHaveSize
 import io.kotest.matchers.shouldBe
 import me.eater.threedom.dom.Document
 import me.eater.threedom.dom.IDocument
@@ -50,7 +52,7 @@ class PositionTest : StringSpec({
         doc.findAt(-10, 40, 0) shouldHaveSingleElement nodeThree
     }
 
-    "ensure in selection works" {
+    "ensure in region works" {
         val doc: IDocument = Document()
         val node = doc.createNode<PlainNode>()
         val nodeTwo = doc.createNode<PlainNode>()
@@ -70,4 +72,25 @@ class PositionTest : StringSpec({
         resultTwo shouldHaveSize 2
         resultTwo shouldContainExactly listOf(nodeTwo, nodeThree)
     }
+
+    "ensure in range 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.findInRange(Vector3d(0, 0, 0), 20).toList()
+        result shouldHaveSize 1
+        result shouldContainExactly listOf(node)
+
+        val resultTwo = doc.findInRange(Vector3d(0, 20, 10), 20).toList()
+        resultTwo shouldHaveSize 2
+        resultTwo shouldContainExactly listOf(nodeTwo, nodeThree)
+    }
 })