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 8f54b12..4036c63 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt +++ b/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) } 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 1f8c0b0..b06aa0d 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/dom/IDocument.kt +++ b/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> fun rebalance() + fun findInRegion(pointA: Vector3dc, pointB: Vector3dc): Sequence> } inline fun > IDocument.createNode() = createNode(T::class) diff --git a/threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeClassListUpdate.kt b/threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeClassListUpdate.kt index 6721704..ead11c6 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeClassListUpdate.kt +++ b/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, override val node: INode<*>): NodeClassListUpdate() - class Added(val classNames: Set, override val node: INode<*>): NodeClassListUpdate() + @EventName("NodeClassesRemoved") + class Removed(val classNames: Set, override val node: INode<*>) : NodeClassListUpdate() + + @EventName("NodeClassesAdded") + class Added(val classNames: Set, override val node: INode<*>) : NodeClassListUpdate() } diff --git a/threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeIDUpdate.kt b/threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeIDUpdate.kt index a2a5b0f..ba60cf4 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/dom/event/NodeIDUpdate.kt +++ b/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?) 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 4614870..0504dfa 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/utils/KDTree.kt +++ b/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() - data class Node( + class Node( val vertex: Vector3dc = Vector3d(0, 0, 0), val nodeIds: MutableSet = mutableSetOf(), val depth: Long = 0, - var left: Node? = null, - var right: Node? = null + val branches: Array = 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 { + var current: Node? = this@Node + val selection = mutableListOf() + 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>> = + 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) diff --git a/threedom/src/main/kotlin/me/eater/threedom/utils/joml/JOML.kt b/threedom/src/main/kotlin/me/eater/threedom/utils/joml/JOML.kt index aff1b41..0784a1c 100644 --- a/threedom/src/main/kotlin/me/eater/threedom/utils/joml/JOML.kt +++ b/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 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 aae2f11..caeaccb 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,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() + 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.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) + } })