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 4036c63..6691ccb 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt +++ b/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) } 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 b06aa0d..7158ad5 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt +++ b/threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt @@ -15,6 +15,7 @@ interface IDocument : EventDispatcher, INodeContainer { fun findAt(vec: Vector3dc): Collection> fun rebalance() fun findInRegion(pointA: Vector3dc, pointB: Vector3dc): Sequence> + fun findInRange(origin: Vector3dc, range: Number): Sequence> } inline fun > 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 0504dfa..6371490 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/utils/KDTree.kt +++ b/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 { + fun findInRange(origin: Vector3dc, range: Number): Sequence { + 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 { var current: Node? = this@Node val selection = mutableListOf() 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) } 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 caeaccb..27751cd 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 @@ -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() val nodeTwo = doc.createNode() @@ -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() + val nodeTwo = doc.createNode() + val nodeThree = doc.createNode() + 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) + } })