Add k-d tree

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

@@ -3,8 +3,11 @@ package me.eater.threedom.dom
 import me.eater.threedom.dom.event.DOMTreeUpdate
 import me.eater.threedom.dom.event.NodeClassListUpdate
 import me.eater.threedom.dom.event.NodeIDUpdate
+import me.eater.threedom.dom.event.NodeModelUpdate
 import me.eater.threedom.event.Event
 import me.eater.threedom.generated.EventNames
+import me.eater.threedom.utils.KDTree
+import org.joml.Vector3dc
 import kotlin.reflect.KClass
 
 class Document : IDocument {
@@ -18,6 +21,10 @@ class Document : IDocument {
             addNodeToSearch(event.child)
         }
 
+        on<DOMTreeUpdate.Move> { (event) ->
+            updateKDTreeForNode(event.child)
+        }
+
         on<NodeIDUpdate> { (event) ->
             event.old?.let {
                 byId.remove(it, event.node.nodeId)
@@ -45,12 +52,27 @@ class Document : IDocument {
                 }
             }
         }
+
+        on<NodeModelUpdate> { (event) ->
+            updateKDTreeForNode(event.node)
+        }
     }
+    private val kdTree = KDTree(this)
 
     private val allNodes: MutableMap<Long, INode<*>> = mutableMapOf()
     private val byId: MutableMap<String, Long> = mutableMapOf()
     private val byClass: MutableMap<String, MutableSet<Long>> = mutableMapOf()
 
+    fun updateKDTreeForNode(node: INode<*>) {
+        node.updateAbsolute()
+        kdTree.update(node)
+
+        for (child in node.recursiveIterator()) {
+            child.updateAbsolute()
+            kdTree.update(child)
+        }
+    }
+
     override fun <T> addEventListener(eventName: String, refNode: INode<*>, block: (Event<T>) -> Unit) {
         @Suppress("UNCHECKED_CAST")
         eventTree.addEventListener(eventName, refNode, block as (Event<*>) -> Unit)
@@ -97,6 +119,8 @@ class Document : IDocument {
         for (className in node.classList) {
             byClass.getOrPut(className, ::mutableSetOf).add(node.nodeId)
         }
+
+        kdTree.add(node)
     }
 
     private fun removeNodeFromSearch(node: INode<*>) {
@@ -111,6 +135,8 @@ class Document : IDocument {
                 byClass.remove(className)
             }
         }
+
+        kdTree.remove(node)
     }
 
     inline fun <reified T> on(topLevel: Boolean = false, noinline block: (Event<T>) -> Unit) = if (topLevel)
@@ -120,6 +146,10 @@ class Document : IDocument {
 
     override fun <T : INode<T>> createNode(nodeType: KClass<T>): T =
         nodeType.java.getConstructor(IDocument::class.java).newInstance(this)
+
+    override fun getNodeByNodeId(nodeId: Long): INode<*>? = this.allNodes[nodeId]
+
+    override fun findAt(vec: Vector3dc) = kdTree.find(vec)
 }
 
 

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

@@ -2,12 +2,17 @@ package me.eater.threedom.dom
 
 import me.eater.threedom.event.Event
 import me.eater.threedom.event.EventDispatcher
+import me.eater.threedom.utils.joml.Vector3d
+import org.joml.Vector3dc
 import kotlin.reflect.KClass
 
 interface IDocument : EventDispatcher, INodeContainer {
     fun <T : INode<T>> createNode(nodeType: KClass<T>): T
     fun deleteNode(refNode: INode<*>)
     fun <T> addEventListener(eventName: String, refNode: INode<*>, block: (Event<T>) -> Unit)
+    fun getNodeByNodeId(nodeId: Long): INode<*>?
+    fun findAt(x: Number, y: Number, z: Number) = findAt(Vector3d(x, y, z))
+    fun findAt(vec: Vector3dc): Collection<INode<*>>
 }
 
 inline fun <reified T : INode<T>> IDocument.createNode() = createNode(T::class)

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

@@ -1,7 +1,8 @@
 package me.eater.threedom.dom
 
+import me.eater.threedom.utils.joml.times
 import org.joml.Matrix4d
-import org.joml.Vector3d
+import org.joml.Matrix4dc
 import java.util.concurrent.atomic.AtomicLong
 
 interface INode<T : INode<T>> : Comparable<INode<*>>, INodeContainer {
@@ -10,10 +11,10 @@ interface INode<T : INode<T>> : Comparable<INode<*>>, INodeContainer {
     val nodeId: Long
     val parentNode: INode<*>?
     val document: IDocument?
-    val absolute: Matrix4d
+    val absolute: Matrix4dc
-        get() = (parentNode?.absolute ?: Matrix4d()).mul(model)
+        get() = (parentNode?.absolute ?: Matrix4d()) * model
 
-    var model: Matrix4d
+    var model: Matrix4dc
 
     fun clone(deep: Boolean): T
     fun updateParent(refNode: INode<*>?): Boolean
@@ -33,6 +34,7 @@ interface INode<T : INode<T>> : Comparable<INode<*>>, INodeContainer {
     }
 
     override fun compareTo(other: INode<*>): Int = this.nodeId.compareTo(other.nodeId)
+    fun updateAbsolute()
 
     companion object {
         private val atomicNodeId = AtomicLong(0)

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

@@ -9,4 +9,17 @@ interface INodeContainer : INodeQueryCapable {
     fun sequence(): Sequence<INode<*>>
 
     operator fun iterator(): Iterator<INode<*>>
+    fun recursiveIterator(): Iterator<INode<*>> = sequence<INode<*>> {
+        val iterators = mutableListOf<Iterator<INode<*>>>()
+        var current: Iterator<INode<*>>? = iterator()
+        while (current != null) {
+
+            for (node in current) {
+                yield(node)
+                iterators.add(node.iterator())
+            }
+
+            current = iterators.firstOrNull()?.apply { iterators.removeAt(0) }
+        }
+    }.iterator()
 }

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

@@ -3,16 +3,21 @@ package me.eater.threedom.dom
 import me.eater.threedom.dom.event.DOMTreeUpdate
 import me.eater.threedom.dom.event.NodeClassListUpdate
 import me.eater.threedom.dom.event.NodeIDUpdate
+import me.eater.threedom.dom.event.NodeModelUpdate
 import me.eater.threedom.event.Event
 import me.eater.threedom.event.EventListener
 import me.eater.threedom.event.trigger
 import me.eater.threedom.utils.ObservableSet
+import me.eater.threedom.utils.joml.mutable
+import me.eater.threedom.utils.joml.times
 import org.joml.Matrix4d
+import org.joml.Matrix4dc
 
 abstract class Node<T : INode<T>>(document: IDocument?) : INode<T>, EventListener {
     override var document: IDocument? = document
         protected set
     override val nodeId = INode.getNextNodeId()
+
     private val nodes: MutableSet<INode<*>> = mutableSetOf()
     override var parentNode: INode<*>? = null
         protected set
@@ -27,6 +32,9 @@ abstract class Node<T : INode<T>>(document: IDocument?) : INode<T>, EventListene
             trigger(NodeIDUpdate(this, old, value))
         }
 
+    override var absolute: Matrix4dc = Matrix4d()
+        protected set
+
     override val classList: MutableSet<String> = ObservableSet {
         when (it.action) {
             ObservableSet.Action.Removed -> trigger(NodeClassListUpdate.Removed(it.elements, this))
@@ -40,7 +48,18 @@ abstract class Node<T : INode<T>>(document: IDocument?) : INode<T>, EventListene
         }
     }
 
-    override var model: Matrix4d = Matrix4d()
+    private var _model: Matrix4d = Matrix4d()
+    override var model: Matrix4dc
+        get() = _model
+        set(value) {
+            val old = Matrix4d(_model)
+            _model = value.mutable()
+            trigger(NodeModelUpdate(this, old))
+        }
+
+    fun model(block: Matrix4d.() -> Matrix4dc) {
+        this.model = block(this._model)
+    }
 
     var children: List<INode<*>> = nodes.toList()
 
@@ -174,4 +193,8 @@ abstract class Node<T : INode<T>>(document: IDocument?) : INode<T>, EventListene
         emptySequence()
     else
         document?.getNodesByClassName(className)?.filter { it.hasParent(this) } ?: emptySequence()
+
+    override fun updateAbsolute() {
+        this.absolute = (this.parentNode?.absolute ?: Matrix4d()) * model
+    }
 }

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

@@ -0,0 +1,8 @@
+package me.eater.threedom.dom.event
+
+import me.eater.threedom.dom.INode
+import me.eater.threedom.kapt.EventName
+import org.joml.Matrix4dc
+
+@EventName("NodeModelUpdate")
+data class NodeModelUpdate(val node: INode<*>, val oldModel: Matrix4dc)

threedom/src/main/kotlin/me/eater/threedom/event/EventListener.kt

@@ -1,7 +1,9 @@
 package me.eater.threedom.event
 
+import me.eater.threedom.generated.EventNames
+
 interface EventListener {
     fun <T> on(eventName: String, block: (Event<T>) -> Unit)
 }
 
-inline fun <reified T> EventListener.on(noinline block: (Event<T>) -> Unit) = on(T::class.java.name, block)
+inline fun <reified T> EventListener.on(noinline block: (Event<T>) -> Unit) = on(EventNames.getEventName<T>(), block)

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

@@ -0,0 +1,122 @@
+package me.eater.threedom.utils
+
+import me.eater.threedom.dom.IDocument
+import me.eater.threedom.dom.INode
+import me.eater.threedom.utils.joml.Vector3d
+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))) {
+    private val nodeLocMap = mutableMapOf<Long, Vector3dc>()
+
+    data 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 axis: Int
+            get() = (depth % 3).toInt()
+
+        fun add(translation: Vector3dc, nodeId: Long) {
+            var current = this
+
+            while (true) {
+                if (translation == current.vertex) {
+                    current.nodeIds.add(nodeId)
+                    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!!
+                    }
+                } else {
+                    if (current.right == null) {
+                        current.right = Node(translation, mutableSetOf(nodeId), current.depth + 1)
+                        return
+                    } else {
+                        current = current.right!!
+                    }
+                }
+            }
+        }
+
+        fun remove(translation: Vector3dc, nodeId: Long) {
+            find(translation)?.nodeIds?.remove(nodeId)
+        }
+
+        fun find(translation: Vector3dc): Node? {
+            var current: Node? = this
+            while (current != null) {
+                if (translation == current.vertex) {
+                    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
+                    }
+            }
+
+            return null
+        }
+
+        companion object {
+            @Suppress("UNCHECKED_CAST")
+            fun create(nodes: Collection<INode<*>>, depth: Long = 0): Node =
+                create(nodes.groupBy({ it.absolute.getTranslation() }) { it.nodeId } as Map<Vector3dc, Collection<Long>>,
+                    depth)
+
+            fun create(nodes: Map<Vector3dc, Collection<Long>>, depth: Long = 0): Node {
+                if (nodes.isEmpty()) {
+                    return Node()
+                }
+
+                if (nodes.size == 1) {
+                    val (loc, onlyNodes) = nodes.entries.first()
+                    return Node(loc, onlyNodes.toMutableSet())
+                }
+
+                val axis: Int = (depth % 3).toInt()
+                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 {
+                    nodes.filterKeys(it::contains)
+                }?.let { create(it, depth + 1) }
+                val right = sorted.slice(median + 1..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)
+            }
+        }
+    }
+
+    constructor(document: IDocument, nodes: Collection<INode<*>>) : this(document, Node.create(nodes))
+
+    fun add(node: INode<*>) {
+        val vec = node.absolute.getTranslation()
+        nodeLocMap[node.nodeId] = vec
+        root.add(vec, node.nodeId)
+    }
+
+    fun remove(node: INode<*>) {
+        root.remove(node.absolute.getTranslation(), node.nodeId)
+        nodeLocMap.remove(node.nodeId)
+    }
+
+    fun find(vertex: Vector3dc) = root.find(vertex)?.nodeIds?.mapNotNull(document::getNodeByNodeId) ?: emptyList()
+
+    fun update(node: INode<*>) {
+        nodeLocMap[node.nodeId]?.let { root.remove(it, node.nodeId) }
+        add(node)
+    }
+}

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

@@ -1,45 +0,0 @@
-package me.eater.threedom.utils
-
-class MutableOrderedSetListIterator<T>(private val collection: OrderedSet<T>, private var index: Int = 0) :
-    MutableListIterator<T> {
-
-    override fun hasPrevious() = collection.size > (index - 1) && index > 1
-
-    override fun nextIndex() = index + 1
-    override fun previous(): T {
-        if (!hasPrevious()) {
-            throw NoSuchElementException()
-        }
-
-        index = previousIndex()
-        return collection[index]
-    }
-
-    override fun previousIndex() = index - 1
-    override fun add(element: T) {
-        collection.add(index, element)
-    }
-
-    override fun hasNext() = nextIndex() < collection.size
-
-    override fun next(): T {
-        if (!hasNext()) {
-            throw NoSuchElementException()
-        }
-
-        index = nextIndex()
-        return collection[index]
-    }
-
-    override fun remove() {
-        collection.removeAt(index)
-
-        if (index > 0) {
-            index -= 1
-        }
-    }
-
-    override fun set(element: T) {
-        collection[index] = element
-    }
-}

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

@@ -1,95 +0,0 @@
-package me.eater.threedom.utils
-
-import java.util.*
-
-class OrderedSet<T>() : MutableSet<T>, MutableList<T> {
-    override val size
-        get() = items.size
-
-    private val set: MutableSet<T> = mutableSetOf()
-    private val items: MutableList<T> = mutableListOf()
-
-    constructor(elements: Collection<T>) : this() {
-        addAll(elements)
-    }
-
-    override fun contains(element: T) = set.contains(element)
-    override fun containsAll(elements: Collection<T>) = set.containsAll(elements)
-    override fun isEmpty() = items.isEmpty()
-    override fun iterator() = items.iterator()
-    override operator fun get(index: Int) = items[index]
-    override fun spliterator(): Spliterator<T> = set.spliterator()
-    override fun indexOf(element: T): Int = items.indexOf(element)
-    override fun lastIndexOf(element: T): Int = indexOf(element)
-    override fun listIterator(): MutableOrderedSetListIterator<T> = MutableOrderedSetListIterator(this)
-    override fun listIterator(index: Int): MutableOrderedSetListIterator<T> = MutableOrderedSetListIterator(this, index)
-    override fun subList(fromIndex: Int, toIndex: Int): OrderedSet<T> = OrderedSet(items.subList(fromIndex, toIndex))
-
-    override fun add(element: T): Boolean {
-        if (set.add(element)) {
-            items.add(element)
-            return true
-        }
-
-        return false
-    }
-
-    override fun addAll(elements: Collection<T>): Boolean {
-        return elements.map(::add).any()
-    }
-
-    override fun clear() {
-        set.clear()
-        items.clear()
-    }
-
-    override fun remove(element: T): Boolean {
-        if (set.remove(element)) {
-            items.remove(element)
-            return true
-        }
-
-        return false
-    }
-
-    override fun removeAll(elements: Collection<T>): Boolean {
-        return elements.map(::remove).any()
-    }
-
-    override fun retainAll(elements: Collection<T>): Boolean {
-        return set.toSet().map {
-            if (elements.contains(it)) {
-                false
-            } else {
-                remove(it)
-            }
-        }.any()
-    }
-
-    override fun add(index: Int, element: T) {
-        if (set.add(element)) {
-            items.add(index, element)
-        }
-    }
-
-    override fun addAll(index: Int, elements: Collection<T>): Boolean {
-        return items.addAll(index, elements.filter(set::add))
-    }
-
-    override fun removeAt(index: Int): T {
-        set.remove(items[index])
-        return items.removeAt(index)
-    }
-
-    override fun set(index: Int, element: T): T {
-        if (!set.add(element)) {
-            return element
-        }
-
-        val old = items[index]
-        items[index] = element
-        remove(old)
-        set.add(element)
-        return old
-    }
-}

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

@@ -1,12 +1,29 @@
 package me.eater.threedom.utils.joml
 
 import org.joml.Matrix4d
+import org.joml.Matrix4dc
 import org.joml.Vector3d
+import org.joml.Vector3dc
 
-fun <T : Number> Matrix4d.setTranslation(x: T, y: T, z: T): Matrix4d =
+fun <T : Number> Matrix4dc.setTranslation(x: T, y: T, z: T): Matrix4d =
-    setTranslation(x.toDouble(), y.toDouble(), z.toDouble())
+    Matrix4d(this).setTranslation(x.toDouble(), y.toDouble(), z.toDouble())
 
-fun Matrix4d.getTranslation(): Vector3d = getTranslation(Vector3d())
+fun Matrix4dc.getTranslation(): Vector3d = getTranslation(Vector3d())
+fun Matrix4dc.mutable(): Matrix4d = if (this is Matrix4d) this else Matrix4d(this)
+
+operator fun Matrix4dc.times(rhs: Matrix4dc) = mul(rhs, Matrix4d())
+operator fun Matrix4d.times(rhs: Matrix4dc) = mul(rhs)
 
 @Suppress("FunctionName")
-fun <T : Number> Vector3d(x: T, y: T, z: T) = Vector3d(x.toDouble(), y.toDouble(), z.toDouble())
+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) {
+        val c = this[i].compareTo(rhs[i])
+        if (c != 0) {
+            return c
+        }
+    }
+
+    return 0
+}

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

@@ -1,8 +1,10 @@
 package me.eater.test.threedom.dom
 
 import io.kotest.core.spec.style.StringSpec
+import io.kotest.matchers.collections.shouldHaveSingleElement
 import io.kotest.matchers.shouldBe
 import me.eater.threedom.dom.Document
+import me.eater.threedom.dom.IDocument
 import me.eater.threedom.dom.PlainNode
 import me.eater.threedom.dom.createNode
 import me.eater.threedom.utils.joml.Vector3d
@@ -11,16 +13,36 @@ import me.eater.threedom.utils.joml.setTranslation
 
 class PositionTest : StringSpec({
     "ensure positioning works" {
-        val doc = Document()
+        val doc: IDocument = Document()
         val node = doc.createNode<PlainNode>()
         doc.addNode(node)
-        node.model.setTranslation(10, 0, 10)
+        node.model { setTranslation(10, 0, 10) }
         node.absolute.getTranslation() shouldBe Vector3d(10, 0, 10)
         val nodeTwo = doc.createNode<PlainNode>()
         node.addNode(nodeTwo)
-        nodeTwo.model.setTranslation(-10, 20, 0)
+        nodeTwo.model { setTranslation(-10, 20, 0) }
         nodeTwo.absolute.getTranslation() shouldBe Vector3d(0, 20, 10)
         doc.addNode(nodeTwo)
         nodeTwo.absolute.getTranslation() shouldBe Vector3d(-10, 20, 0)
     }
+
+    "ensure position search tree 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) }
+
+        doc.findAt(10, 0, 10) shouldHaveSingleElement node
+        doc.findAt(0, 20, 10) shouldHaveSingleElement nodeTwo
+        doc.findAt(0, 40, 10) shouldHaveSingleElement nodeThree
+        doc.addNode(nodeTwo)
+        doc.findAt(-10, 20, 0) shouldHaveSingleElement nodeTwo
+        doc.findAt(-10, 40, 0) shouldHaveSingleElement nodeThree
+    }
 })