From c839c64cea9426d90af1a9c36bd8e0d55a33f4e6 Mon Sep 17 00:00:00 2001 From: "Li, Zhen" Date: Thu, 30 Nov 2023 13:32:02 -0500 Subject: [PATCH 1/8] Add Buffered KDTree --- .../ForceSimulation/Forces/CenterForce.swift | 2 +- Sources/ForceSimulation/KDTree/KDTree.swift | 192 +++++++++++++++++- Sources/ForceSimulation/UnsafeArray.swift | 14 ++ 3 files changed, 201 insertions(+), 7 deletions(-) diff --git a/Sources/ForceSimulation/Forces/CenterForce.swift b/Sources/ForceSimulation/Forces/CenterForce.swift index c3328df..a9efeae 100644 --- a/Sources/ForceSimulation/Forces/CenterForce.swift +++ b/Sources/ForceSimulation/Forces/CenterForce.swift @@ -6,7 +6,7 @@ extension Kinetics { /// See [Collide Force - D3](https://d3js.org/d3-force/collide). public struct CenterForce: ForceProtocol { - @usableFromInline var kinetics: Kinetics! = nil + @usableFromInline var kinetics: Kinetics! = nil @inlinable public func apply() { diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index 742b8d5..1acb591 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -1,3 +1,183 @@ +public struct KDTreeNode +where + Vector: SimulatableVector & L2NormCalculatable, + Delegate: KDTreeDelegate +{ + public struct NodeIndex { + var index: Int + var next: UnsafeMutablePointer? + } + public var nodeIndices: NodeIndex? + public var childrenBufferPointer: UnsafeMutablePointer>? + public var delegate: Delegate + + @inlinable + init( + nodeIndices: NodeIndex?, + childrenBufferPointer: UnsafeMutablePointer>?, + delegate: consuming Delegate + ) { + self.childrenBufferPointer = childrenBufferPointer + self.nodeIndices = nodeIndices + self.delegate = consume delegate + } +} + +public struct KDTreeRoot +where + Vector: SimulatableVector & L2NormCalculatable, + Delegate: KDTreeDelegate +{ + public typealias Box = KDBox + public typealias TreeNode = KDTreeNode + public let root: UnsafeMutablePointer + public let treeNodeBuffer: UnsafeArray + public var box: Box + + /// + public var validCount: Int = 0 + + @inlinable + var clusterDistanceSquared: Vector.Scalar { + return Vector.clusterDistanceSquared + } + + @inlinable + public init( + box: Box, + rootDelegate: @autoclosure () -> Delegate + ) { + let maxBufferCount = 1 << Vector.scalarCount + let zeroNode: TreeNode = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate() + ) + treeNodeBuffer = .createBuffer( + withHeader: maxBufferCount, + count: maxBufferCount, + initialValue: zeroNode + ) + root = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } + + root.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate() + ) + + self.box = box + self.validCount = 1 + } + + @inlinable + public mutating func reset( + rootDelegate: @autoclosure () -> Delegate + ) { + root.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate() + ) + self.validCount = 1 + } + + + @inlinable + public mutating func add(index: Int, at point: Vector) { + assert(validCount > 0) + } + + @inlinable + internal func addWithoutCover(index: Int, at point: Vector) { + + } + + @inlinable + internal mutating func cover(point: Vector) { + if box.contains(point) { return } + + repeat { + let direction = self.getIndexInChildren(point, relativeTo: box.p0) + self.expand(towards: direction) + } while !box.contains(point) + } + + @inlinable + internal mutating func expand(towards direction: Int) { + let nailedDirection = (Self.directionCount - 1) - direction + let nailedCorner = box.getCorner(of: nailedDirection) + let _corner = box.getCorner(of: direction) + let expandedCorner = (_corner + _corner) - nailedCorner + let newRootBox = Box(nailedCorner, expandedCorner) + + let _rootValue = self.root.pointee + + + // spawn the delegate with the same internal values + // for the children, use implicit copy of spawned + let spawned = _rootValue.delegate.spawn() + + + let newChildrenPointer = self.root + validCount + + for j in 0..> i) & 0b1 + // TODO: use simd mask + if isOnTheHigherRange != 0 { + __box.p0[i] = _corner[i] + } else { + __box.p1[i] = _corner[i] + } + } + + self.treeNodeBuffer[validCount + j] = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: j != nailedDirection ? _rootValue.delegate : spawned + ) + } + self.validCount += Self.directionCount + + self.root.pointee = .init(nodeIndices: nil, childrenBufferPointer: newChildrenPointer, delegate: _rootValue.delegate) + } + + @inlinable + internal mutating func resizeBuffer(to capacity: UInt) { + // TODO + } + + @inlinable + static var directionCount: Int { 1 << Vector.scalarCount } + + @inlinable + internal func deinitializeBuffer() { + _ = treeNodeBuffer.withUnsafeMutablePointerToElements { + $0.deinitialize(count: Self.directionCount) + } + } + + /// Get the index of the child that contains the point. + /// + /// **Complexity**: `O(n*(2^n))`, where `n` is the dimension of the vector. + @inlinable + func getIndexInChildren(_ point: Vector, relativeTo originalPoint: Vector) -> Int { + var index = 0 + + let mask = point .>= originalPoint + + for i in 0..) -> Bool) { + @inlinable public mutating func visit( + shouldVisitChildren: (inout KDTree) -> Bool + ) { if shouldVisitChildren(&self) && children != nil { // this is an internal node for i in children!.indices { @@ -314,7 +496,6 @@ extension KDTree { } - // public struct KDTreeRoot // where // Vector: SimulatableVector & L2NormCalculatable, @@ -332,10 +513,9 @@ extension KDTree { // self.propertyBuffer = propertyBuffer // } - // @inlinable // public mutating func add(_ nodeIndex: Int, at point: Vector) { // root.cover(point) // root.addWithoutCover(nodeIndex, at: point) // } -// } \ No newline at end of file +// } diff --git a/Sources/ForceSimulation/UnsafeArray.swift b/Sources/ForceSimulation/UnsafeArray.swift index bbb5cca..17a55ff 100644 --- a/Sources/ForceSimulation/UnsafeArray.swift +++ b/Sources/ForceSimulation/UnsafeArray.swift @@ -11,6 +11,20 @@ public final class UnsafeArray: ManagedBuffer { } return unsafeDowncast(buffer, to: UnsafeArray.self) } + + @inlinable + class func createUninitializedBuffer( + count: Int + ) -> UnsafeArray { + let buffer = self.create(minimumCapacity: count) { _ in count } + // buffer.withUnsafeMutablePointerToElements { + // $0.initialize(repeating: Element(), count: count) + // } + return unsafeDowncast(buffer, to: UnsafeArray.self) + } + + + @inlinable var count: Int { return header From ec52135c491dde213e21481df52f82f2a5c615db Mon Sep 17 00:00:00 2001 From: "Li, Zhen" Date: Fri, 1 Dec 2023 08:06:14 -0500 Subject: [PATCH 2/8] Add BufferedKDTree --- Sources/ForceSimulation/KDTree/KDTree.swift | 191 +++++++++++++++++--- Tests/KDTreeTests/BufferedKDTreeTests.swift | 11 ++ 2 files changed, 173 insertions(+), 29 deletions(-) create mode 100644 Tests/KDTreeTests/BufferedKDTreeTests.swift diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index 1acb591..e595e38 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -4,10 +4,43 @@ where Delegate: KDTreeDelegate { public struct NodeIndex { + + @usableFromInline var index: Int + + @usableFromInline var next: UnsafeMutablePointer? + + @inlinable + internal init( + nodeIndex: Int + ) { + self.index = nodeIndex + self.next = nil + } + + @inlinable + internal mutating func append(nodeIndex: Int) { + if let next { + next.pointee.append(nodeIndex: nodeIndex) + } + else { + next = .allocate(capacity: 1) + next!.pointee = .init(nodeIndex: nodeIndex) + } + } + + @inlinable + internal func deinitialize() { + if let next { + next.pointee.deinitialize() + next.deallocate() + } + } } public var nodeIndices: NodeIndex? + public var box: KDBox + public var nodePosition: Vector public var childrenBufferPointer: UnsafeMutablePointer>? public var delegate: Delegate @@ -15,15 +48,20 @@ where init( nodeIndices: NodeIndex?, childrenBufferPointer: UnsafeMutablePointer>?, - delegate: consuming Delegate + delegate: consuming Delegate, + box: consuming KDBox ) { self.childrenBufferPointer = childrenBufferPointer self.nodeIndices = nodeIndices self.delegate = consume delegate + self.box = consume box + self.nodePosition = .zero } + + } -public struct KDTreeRoot +public struct BufferedKDTree: ~Copyable where Vector: SimulatableVector & L2NormCalculatable, Delegate: KDTreeDelegate @@ -32,7 +70,7 @@ where public typealias TreeNode = KDTreeNode public let root: UnsafeMutablePointer public let treeNodeBuffer: UnsafeArray - public var box: Box + public var rootBox: Box /// public var validCount: Int = 0 @@ -44,29 +82,31 @@ where @inlinable public init( - box: Box, + rootBox: Box, rootDelegate: @autoclosure () -> Delegate ) { let maxBufferCount = 1 << Vector.scalarCount let zeroNode: TreeNode = .init( - nodeIndices: nil, + nodeIndices: nil, childrenBufferPointer: nil, - delegate: rootDelegate() + delegate: rootDelegate(), + box: rootBox ) treeNodeBuffer = .createBuffer( - withHeader: maxBufferCount, - count: maxBufferCount, + withHeader: maxBufferCount, + count: maxBufferCount, initialValue: zeroNode ) root = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } root.pointee = .init( - nodeIndices: nil, + nodeIndices: nil, childrenBufferPointer: nil, - delegate: rootDelegate() + delegate: rootDelegate(), + box: rootBox ) - self.box = box + self.rootBox = rootBox self.validCount = 1 } @@ -75,55 +115,142 @@ where rootDelegate: @autoclosure () -> Delegate ) { root.pointee = .init( - nodeIndices: nil, + nodeIndices: nil, childrenBufferPointer: nil, - delegate: rootDelegate() + delegate: rootDelegate(), + box: rootBox ) self.validCount = 1 } - @inlinable - public mutating func add(index: Int, at point: Vector) { + public mutating func add( + nodeIndex: Int, + at point: Vector + ) { assert(validCount > 0) + cover(point: point) + addWithoutCover( + onTreeNode: root, + nodeOf: nodeIndex, + at: point + ) } @inlinable - internal func addWithoutCover(index: Int, at point: Vector) { + internal mutating func addWithoutCover( + onTreeNode treeNode: UnsafeMutablePointer, + nodeOf nodeIndex: Int, + at point: Vector + ) { + + defer { + treeNode.pointee.delegate.didAddNode(nodeIndex, at: point) + } + guard treeNode.pointee.childrenBufferPointer != nil else { + if treeNode.pointee.nodeIndices == nil { + treeNode.pointee.nodeIndices = .init(nodeIndex: nodeIndex) + treeNode.pointee.nodePosition = point + return + } else if treeNode.pointee.nodePosition == point + || treeNode.pointee.nodePosition.distanceSquared(to: point) < clusterDistanceSquared + { + treeNode.pointee.nodeIndices!.append(nodeIndex: nodeIndex) + return + } else { + + let spawnedDelegate = treeNode.pointee.delegate.spawn() + let center = treeNode.pointee.box.center + + for j in 0..> i) & 0b1 + // TODO: use simd mask + if isOnTheHigherRange != 0 { + __box.p0[i] = center[i] + } else { + __box.p1[i] = center[i] + } + } + + treeNodeBuffer[validCount+j] = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: spawnedDelegate, + box: __box + ) + } + validCount += Self.directionCount + + if let childrenBufferPointer = treeNode.pointee.childrenBufferPointer { + let direction = getIndexInChildren( + treeNode.pointee.nodePosition, + relativeTo: center + ) + + childrenBufferPointer[direction].nodeIndices = treeNode.pointee.nodeIndices + childrenBufferPointer[direction].nodePosition = treeNode.pointee.nodePosition + childrenBufferPointer[direction].delegate = treeNode.pointee.delegate + treeNode.pointee.nodeIndices = nil + treeNode.pointee.nodePosition = .zero + } + + let directionOfNewNode = getIndexInChildren(point, relativeTo: center) + // spawnedChildren[directionOfNewNode].addWithoutCover(nodeIndex, at: point) + addWithoutCover( + onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, + nodeOf: nodeIndex, + at: point + ) + + + // self.children = spawnedChildren + return + } + } + + let directionOfNewNode = getIndexInChildren(point, relativeTo: treeNode.pointee.box.center) + self.addWithoutCover( + onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, + nodeOf: nodeIndex, + at: point + ) + return } @inlinable internal mutating func cover(point: Vector) { - if box.contains(point) { return } + if rootBox.contains(point) { return } repeat { - let direction = self.getIndexInChildren(point, relativeTo: box.p0) + let direction = self.getIndexInChildren(point, relativeTo: rootBox.p0) self.expand(towards: direction) - } while !box.contains(point) + } while !rootBox.contains(point) } @inlinable internal mutating func expand(towards direction: Int) { let nailedDirection = (Self.directionCount - 1) - direction - let nailedCorner = box.getCorner(of: nailedDirection) - let _corner = box.getCorner(of: direction) + let nailedCorner = rootBox.getCorner(of: nailedDirection) + let _corner = rootBox.getCorner(of: direction) let expandedCorner = (_corner + _corner) - nailedCorner let newRootBox = Box(nailedCorner, expandedCorner) let _rootValue = self.root.pointee - // spawn the delegate with the same internal values // for the children, use implicit copy of spawned let spawned = _rootValue.delegate.spawn() - let newChildrenPointer = self.root + validCount for j in 0..> i) & 0b1 // TODO: use simd mask @@ -133,16 +260,22 @@ where __box.p1[i] = _corner[i] } } - + self.treeNodeBuffer[validCount + j] = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: j != nailedDirection ? _rootValue.delegate : spawned + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: j != nailedDirection ? _rootValue.delegate : spawned, + box: __box ) } self.validCount += Self.directionCount - - self.root.pointee = .init(nodeIndices: nil, childrenBufferPointer: newChildrenPointer, delegate: _rootValue.delegate) + + self.root.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: newChildrenPointer, + delegate: _rootValue.delegate, + box: newRootBox + ) } @inlinable diff --git a/Tests/KDTreeTests/BufferedKDTreeTests.swift b/Tests/KDTreeTests/BufferedKDTreeTests.swift new file mode 100644 index 0000000..cd236d4 --- /dev/null +++ b/Tests/KDTreeTests/BufferedKDTreeTests.swift @@ -0,0 +1,11 @@ +import XCTest +@testable import ForceSimulation + + +class BufferedKDTreeTests: XCTestCase { + func testEmptyTree() { + let tree = BufferedKDTree(box: KDBox([0, 0], [1, 1])) + assert(tree.debugDescription ~= "[]") + assert(tree.delegate.count == 0) + } +} \ No newline at end of file From 1978cbed3efe5427ed40358c4e4e0c27e0a4af8c Mon Sep 17 00:00:00 2001 From: "Li, Zhen" Date: Fri, 1 Dec 2023 08:36:47 -0500 Subject: [PATCH 3/8] Adding tests --- Package.swift | 8 +-- Sources/ForceSimulation/KDTree/KDTree.swift | 18 ++--- Tests/KDTreeTests/BufferedKDTreeTests.swift | 74 +++++++++++++++++++-- 3 files changed, 83 insertions(+), 17 deletions(-) diff --git a/Package.swift b/Package.swift index 462c655..76a746b 100644 --- a/Package.swift +++ b/Package.swift @@ -45,10 +45,10 @@ let package = Package( // swiftSettings: [.unsafeFlags(["-Xfrontend", "-disable-availability-checking"])] ), - // .testTarget( - // name: "NDTreeTests", - // dependencies: ["ForceSimulation"] - // ), + .testTarget( + name: "KDTreeTests", + dependencies: ["ForceSimulation"] + ), .testTarget( name: "ForceSimulationTests", diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index e595e38..22ae392 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -70,7 +70,6 @@ where public typealias TreeNode = KDTreeNode public let root: UnsafeMutablePointer public let treeNodeBuffer: UnsafeArray - public var rootBox: Box /// public var validCount: Int = 0 @@ -83,9 +82,10 @@ where @inlinable public init( rootBox: Box, + nodeCapacity: Int, rootDelegate: @autoclosure () -> Delegate ) { - let maxBufferCount = 1 << Vector.scalarCount + let maxBufferCount = (nodeCapacity << Vector.scalarCount) + 1 let zeroNode: TreeNode = .init( nodeIndices: nil, childrenBufferPointer: nil, @@ -106,12 +106,13 @@ where box: rootBox ) - self.rootBox = rootBox +// self.rootBox = rootBox self.validCount = 1 } @inlinable public mutating func reset( + rootBox: Box, rootDelegate: @autoclosure () -> Delegate ) { root.pointee = .init( @@ -182,6 +183,7 @@ where box: __box ) } + treeNode.pointee.childrenBufferPointer = root + validCount validCount += Self.directionCount if let childrenBufferPointer = treeNode.pointee.childrenBufferPointer { @@ -223,19 +225,19 @@ where @inlinable internal mutating func cover(point: Vector) { - if rootBox.contains(point) { return } + if self.root.pointee.box.contains(point) { return } repeat { - let direction = self.getIndexInChildren(point, relativeTo: rootBox.p0) + let direction = self.getIndexInChildren(point, relativeTo: self.root.pointee.box.p0) self.expand(towards: direction) - } while !rootBox.contains(point) + } while !self.root.pointee.box.contains(point) } @inlinable internal mutating func expand(towards direction: Int) { let nailedDirection = (Self.directionCount - 1) - direction - let nailedCorner = rootBox.getCorner(of: nailedDirection) - let _corner = rootBox.getCorner(of: direction) + let nailedCorner = self.root.pointee.box.getCorner(of: nailedDirection) + let _corner = self.root.pointee.box.getCorner(of: direction) let expandedCorner = (_corner + _corner) - nailedCorner let newRootBox = Box(nailedCorner, expandedCorner) diff --git a/Tests/KDTreeTests/BufferedKDTreeTests.swift b/Tests/KDTreeTests/BufferedKDTreeTests.swift index cd236d4..2468779 100644 --- a/Tests/KDTreeTests/BufferedKDTreeTests.swift +++ b/Tests/KDTreeTests/BufferedKDTreeTests.swift @@ -1,11 +1,75 @@ import XCTest @testable import ForceSimulation +struct CountKDTreeDelegate: KDTreeDelegate { + mutating func didAddNode(_ node: Int, at position: SIMD2) { + count += 1 + } + + mutating func didRemoveNode(_ node: Int, at position: SIMD2) { + count -= 1 + } + + + typealias NodeID = Int + + typealias Vector = SIMD2 + + var count = 0 + + func spawn() -> CountKDTreeDelegate { + return .init(count: 0) + } + +} class BufferedKDTreeTests: XCTestCase { - func testEmptyTree() { - let tree = BufferedKDTree(box: KDBox([0, 0], [1, 1])) - assert(tree.debugDescription ~= "[]") - assert(tree.delegate.count == 0) + + private func buildTree( + box: KDBox>, + points: [SIMD2] + ) -> BufferedKDTree, CountKDTreeDelegate> { + var t = BufferedKDTree( + rootBox: box, + nodeCapacity: points.count, + rootDelegate: CountKDTreeDelegate() + ) + for i in points.indices { + t.add(nodeIndex: i, at: points[i]) + } + return t + } + + func testCorner() { + + let t = buildTree(box: .init(p0: [0,0], p1: [1,1]), points: [ + [0,0] + ]) + + XCTAssert(t.root.pointee.nodeIndices!.index == 0) + XCTAssert(t.root.pointee.childrenBufferPointer == nil) + XCTAssert(t.root.pointee.delegate.count == 1) } -} \ No newline at end of file + + func testCorner2() { + let t = buildTree(box: .init(p0: [0,0], p1: [1,1]), points: [ + [1,1] + ]) + + XCTAssert(t.root.pointee.nodeIndices == nil) + XCTAssert(t.root.pointee.delegate.count == 1) + XCTAssert(t.root.pointee.childrenBufferPointer![3].delegate.count == 1) + XCTAssert(t.root.pointee.box.p1 == [2, 2]) + } + + func testRandomTree() { + let randomPoints = (0..<1000).map { _ in + SIMD2([Double.random(in: 0..<100), Double.random(in: 0..<100)]) + } + + let t = buildTree(box: .init(p0: [0,0], p1: [100,100]), points: randomPoints) + XCTAssert(t.root.pointee.delegate.count == randomPoints.count) + } + + +} From 4b87d9a6cecb39d401398a5172c36e9afcb2d4f7 Mon Sep 17 00:00:00 2001 From: "Li, Zhen" Date: Fri, 1 Dec 2023 17:51:52 -0500 Subject: [PATCH 4/8] Add tests --- Sources/ForceSimulation/KDTree/KDTree.swift | 133 ++++++++++++-------- Tests/KDTreeTests/BufferedKDTreeTests.swift | 16 +-- 2 files changed, 89 insertions(+), 60 deletions(-) diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index 22ae392..c5ca407 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -5,7 +5,7 @@ where { public struct NodeIndex { - @usableFromInline + @usableFromInline var index: Int @usableFromInline @@ -23,8 +23,7 @@ where internal mutating func append(nodeIndex: Int) { if let next { next.pointee.append(nodeIndex: nodeIndex) - } - else { + } else { next = .allocate(capacity: 1) next!.pointee = .init(nodeIndex: nodeIndex) } @@ -58,7 +57,6 @@ where self.nodePosition = .zero } - } public struct BufferedKDTree: ~Copyable @@ -68,17 +66,24 @@ where { public typealias Box = KDBox public typealias TreeNode = KDTreeNode - public let root: UnsafeMutablePointer - public let treeNodeBuffer: UnsafeArray - /// - public var validCount: Int = 0 + @usableFromInline + var rootPointer: UnsafeMutablePointer + + @usableFromInline + internal var validCount: Int = 0 + + @usableFromInline + internal var treeNodeBuffer: UnsafeArray @inlinable var clusterDistanceSquared: Vector.Scalar { return Vector.clusterDistanceSquared } + @inlinable + public var root: TreeNode { rootPointer.pointee } + @inlinable public init( rootBox: Box, @@ -97,16 +102,14 @@ where count: maxBufferCount, initialValue: zeroNode ) - root = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } + rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } - root.pointee = .init( + rootPointer.pointee = .init( nodeIndices: nil, childrenBufferPointer: nil, delegate: rootDelegate(), box: rootBox ) - -// self.rootBox = rootBox self.validCount = 1 } @@ -115,7 +118,39 @@ where rootBox: Box, rootDelegate: @autoclosure () -> Delegate ) { - root.pointee = .init( + rootPointer.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate(), + box: rootBox + ) + self.validCount = 1 + } + + @inlinable + public mutating func discardAndResize( + rootBox: Box, + rootDelegate: @autoclosure () -> Delegate, + forNodeCapacity capacity: Int + ) { + let maxBufferCount = (capacity << Vector.scalarCount) + 1 + + let zeroNode: TreeNode = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate(), + box: rootBox + ) + + // should be deinitialized here (ManagedBuffer) + treeNodeBuffer = .createBuffer( + withHeader: maxBufferCount, + count: maxBufferCount, + initialValue: zeroNode + ) + rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } + + rootPointer.pointee = .init( nodeIndices: nil, childrenBufferPointer: nil, delegate: rootDelegate(), @@ -126,14 +161,14 @@ where @inlinable public mutating func add( - nodeIndex: Int, + nodeIndex: Int, at point: Vector ) { assert(validCount > 0) cover(point: point) addWithoutCover( - onTreeNode: root, - nodeOf: nodeIndex, + onTreeNode: rootPointer, + nodeOf: nodeIndex, at: point ) } @@ -144,7 +179,7 @@ where nodeOf nodeIndex: Int, at point: Vector ) { - + defer { treeNode.pointee.delegate.didAddNode(nodeIndex, at: point) } @@ -159,7 +194,7 @@ where treeNode.pointee.nodeIndices!.append(nodeIndex: nodeIndex) return } else { - + let spawnedDelegate = treeNode.pointee.delegate.spawn() let center = treeNode.pointee.box.center @@ -175,20 +210,20 @@ where __box.p1[i] = center[i] } } - - treeNodeBuffer[validCount+j] = .init( - nodeIndices: nil, - childrenBufferPointer: nil, + + treeNodeBuffer[validCount + j] = .init( + nodeIndices: nil, + childrenBufferPointer: nil, delegate: spawnedDelegate, box: __box ) } - treeNode.pointee.childrenBufferPointer = root + validCount + treeNode.pointee.childrenBufferPointer = rootPointer + validCount validCount += Self.directionCount if let childrenBufferPointer = treeNode.pointee.childrenBufferPointer { let direction = getIndexInChildren( - treeNode.pointee.nodePosition, + treeNode.pointee.nodePosition, relativeTo: center ) @@ -202,12 +237,11 @@ where let directionOfNewNode = getIndexInChildren(point, relativeTo: center) // spawnedChildren[directionOfNewNode].addWithoutCover(nodeIndex, at: point) addWithoutCover( - onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, - nodeOf: nodeIndex, + onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, + nodeOf: nodeIndex, at: point ) - // self.children = spawnedChildren return @@ -216,8 +250,8 @@ where let directionOfNewNode = getIndexInChildren(point, relativeTo: treeNode.pointee.box.center) self.addWithoutCover( - onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, - nodeOf: nodeIndex, + onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, + nodeOf: nodeIndex, at: point ) return @@ -225,29 +259,29 @@ where @inlinable internal mutating func cover(point: Vector) { - if self.root.pointee.box.contains(point) { return } + if self.root.box.contains(point) { return } repeat { - let direction = self.getIndexInChildren(point, relativeTo: self.root.pointee.box.p0) + let direction = self.getIndexInChildren(point, relativeTo: self.root.box.p0) self.expand(towards: direction) - } while !self.root.pointee.box.contains(point) + } while !self.root.box.contains(point) } @inlinable internal mutating func expand(towards direction: Int) { let nailedDirection = (Self.directionCount - 1) - direction - let nailedCorner = self.root.pointee.box.getCorner(of: nailedDirection) - let _corner = self.root.pointee.box.getCorner(of: direction) + let nailedCorner = self.root.box.getCorner(of: nailedDirection) + let _corner = self.root.box.getCorner(of: direction) let expandedCorner = (_corner + _corner) - nailedCorner let newRootBox = Box(nailedCorner, expandedCorner) - let _rootValue = self.root.pointee + let _rootValue = self.root // spawn the delegate with the same internal values // for the children, use implicit copy of spawned let spawned = _rootValue.delegate.spawn() - let newChildrenPointer = self.root + validCount + let newChildrenPointer = self.rootPointer + validCount for j in 0..( // covering points: [T], // keyPath: KeyPath, @@ -636,21 +665,21 @@ extension KDTree { // Vector: SimulatableVector & L2NormCalculatable, // Delegate: KDTreeDelegate // { -// public var root: KDTree +// public var rootPointer: KDTree // @usableFromInline let propertyBuffer: UnsafeMutablePointer // @inlinable // public init( -// root: KDTree, +// rootPointer: KDTree, // propertyBuffer: UnsafeMutablePointer // ) { -// self.root = root +// self.rootPointer = rootPointer // self.propertyBuffer = propertyBuffer // } // @inlinable // public mutating func add(_ nodeIndex: Int, at point: Vector) { -// root.cover(point) -// root.addWithoutCover(nodeIndex, at: point) +// rootPointer.cover(point) +// rootPointer.addWithoutCover(nodeIndex, at: point) // } // } diff --git a/Tests/KDTreeTests/BufferedKDTreeTests.swift b/Tests/KDTreeTests/BufferedKDTreeTests.swift index 2468779..652c9c7 100644 --- a/Tests/KDTreeTests/BufferedKDTreeTests.swift +++ b/Tests/KDTreeTests/BufferedKDTreeTests.swift @@ -46,9 +46,9 @@ class BufferedKDTreeTests: XCTestCase { [0,0] ]) - XCTAssert(t.root.pointee.nodeIndices!.index == 0) - XCTAssert(t.root.pointee.childrenBufferPointer == nil) - XCTAssert(t.root.pointee.delegate.count == 1) + XCTAssert(t.root.nodeIndices!.index == 0) + XCTAssert(t.root.childrenBufferPointer == nil) + XCTAssert(t.root.delegate.count == 1) } func testCorner2() { @@ -56,10 +56,10 @@ class BufferedKDTreeTests: XCTestCase { [1,1] ]) - XCTAssert(t.root.pointee.nodeIndices == nil) - XCTAssert(t.root.pointee.delegate.count == 1) - XCTAssert(t.root.pointee.childrenBufferPointer![3].delegate.count == 1) - XCTAssert(t.root.pointee.box.p1 == [2, 2]) + XCTAssert(t.root.nodeIndices == nil) + XCTAssert(t.root.delegate.count == 1) + XCTAssert(t.root.childrenBufferPointer![3].delegate.count == 1) + XCTAssert(t.root.box.p1 == [2, 2]) } func testRandomTree() { @@ -68,7 +68,7 @@ class BufferedKDTreeTests: XCTestCase { } let t = buildTree(box: .init(p0: [0,0], p1: [100,100]), points: randomPoints) - XCTAssert(t.root.pointee.delegate.count == randomPoints.count) + XCTAssert(t.root.delegate.count == randomPoints.count) } From 5c402f110dfafa4e0176da0ae5f016677ebc90c5 Mon Sep 17 00:00:00 2001 From: "Li, Zhen" Date: Fri, 1 Dec 2023 22:35:54 -0500 Subject: [PATCH 5/8] Add tree deinit --- Sources/ForceSimulation/ForceProtocol.swift | 10 +- .../ForceSimulation/Forces/CollideForce.swift | 107 +++++++++--------- .../Forces/ManyBodyForce.swift | 47 ++++++-- Sources/ForceSimulation/KDTree/KDTree.swift | 66 ++++++++++- Sources/ForceSimulation/Simulation.swift | 7 +- .../GKTreeCompareTest.swift | 9 +- .../MiserableGraphTest.swift | 8 +- 7 files changed, 179 insertions(+), 75 deletions(-) diff --git a/Sources/ForceSimulation/ForceProtocol.swift b/Sources/ForceSimulation/ForceProtocol.swift index 18a99b1..7a448a0 100644 --- a/Sources/ForceSimulation/ForceProtocol.swift +++ b/Sources/ForceSimulation/ForceProtocol.swift @@ -12,6 +12,14 @@ public protocol ForceProtocol { /// Bind to a kinetic system that describes the state of all nodes in your simulation. /// This has to be called before `apply` is called. @inlinable mutating func bindKinetics(_ kinetics: Kinetics) + +} + +public extension ForceProtocol { + @inlinable + func dispose() { + + } } public protocol Force2D: ForceProtocol where Vector == SIMD2 {} @@ -67,4 +75,4 @@ extension ForceField { public protocol ForceField2D: ForceField & Force2D { } -public protocol ForceField3D: ForceField & Force3D { } \ No newline at end of file +public protocol ForceField3D: ForceField & Force3D { } diff --git a/Sources/ForceSimulation/Forces/CollideForce.swift b/Sources/ForceSimulation/Forces/CollideForce.swift index b2157d2..31600b4 100644 --- a/Sources/ForceSimulation/Forces/CollideForce.swift +++ b/Sources/ForceSimulation/Forces/CollideForce.swift @@ -68,10 +68,27 @@ extension Kinetics { public mutating func bindKinetics(_ kinetics: Kinetics) { self.kinetics = kinetics self.calculatedRadius = self.radius.calculateUnsafe(for: kinetics.validCount) + self.tree = .allocate(capacity: 1) + self.tree.initialize( + to: + BufferedKDTree( + rootBox: .init( + p0: .init(repeating: 0), + p1: .init(repeating: 1) + ), + nodeCapacity: kinetics.validCount, + rootDelegate: .init(radiusBufferPointer: self.calculatedRadius.mutablePointer) + ) + ) } @usableFromInline var calculatedRadius: UnsafeArray! = nil + + @usableFromInline + internal var tree: + UnsafeMutablePointer>>! = nil + @inlinable public func apply() { assert(self.kinetics != nil, "Kinetics not bound to force") @@ -82,14 +99,20 @@ extension Kinetics { let positionBufferPointer = kinetics.position.mutablePointer let velocityBufferPointer = kinetics.velocity.mutablePointer + let tree = self.tree! + for _ in 0..>( - covering: kinetics.position, - rootDelegate: MaxRadiusNDTreeDelegate( - radiusBufferPointer: calculatedRadius - ) + let coveringBox = KDBox.cover(of: self.kinetics.position) + tree.pointee.reset( + rootBox: coveringBox, + rootDelegate: .init(radiusBufferPointer: calculatedRadius) ) + assert(tree.pointee.validCount == 1) + for p in kinetics.range { +// print("==\(tree.pointee.validCount)") + tree.pointee.add(nodeIndex: p, at: positionBufferPointer[p]) + } for i in kinetics.range { let iOriginalPosition = positionBufferPointer[i] @@ -99,60 +122,44 @@ extension Kinetics { let iPosition = iOriginalPosition + iOriginalVelocity let random = kinetics.randomGenerator - tree.visit { t in + tree.pointee.visit { t in let maxRadiusOfQuad = t.delegate.maxNodeRadius let deltaR = maxRadiusOfQuad + iR - if t.nodePosition != nil { - for j in t.nodeIndices { - // print("\(i)<=>\(j)") - // is leaf, make sure every collision happens once. - guard j > i else { continue } + if let jLinkNode = t.nodeIndices { + let j = jLinkNode.index + // print("\(i)<=>\(j)") + // is leaf, make sure every collision happens once. + guard j > i else { return false } - let jR = calculatedRadius[j] - let jOriginalPosition = positionBufferPointer[j] - let jOriginalVelocity = velocityBufferPointer[j] - var deltaPosition = - iPosition - (jOriginalPosition + jOriginalVelocity) - let l = (deltaPosition).lengthSquared() + let jR = calculatedRadius[j] + let jOriginalPosition = positionBufferPointer[j] + let jOriginalVelocity = velocityBufferPointer[j] + var deltaPosition = + iPosition - (jOriginalPosition + jOriginalVelocity) + let l = (deltaPosition).lengthSquared() - let deltaR = iR + jR - if l < deltaR * deltaR { + let deltaR = iR + jR + if l < deltaR * deltaR { - var l = /*simd_length*/ (deltaPosition.jiggled(by: random)) - .length() - l = (deltaR - l) / l * strength + var l = /*simd_length*/ (deltaPosition.jiggled(by: random)) + .length() + l = (deltaR - l) / l * strength - let jR2 = jR * jR + let jR2 = jR * jR - let k = jR2 / (iR2 + jR2) + let k = jR2 / (iR2 + jR2) - deltaPosition *= l + deltaPosition *= l - velocityBufferPointer[i] += deltaPosition * k - velocityBufferPointer[j] -= deltaPosition * (1 - k) - } + velocityBufferPointer[i] += deltaPosition * k + velocityBufferPointer[j] -= deltaPosition * (1 - k) } + return false } - // TODO: SIMD mask - - // for laneIndex in t.box.p0.indices { - // let _v = t.box.p0[laneIndex] - // if _v > iPosition[laneIndex] + deltaR /* True if no overlap */ { - // return false - // } - // } - - // for laneIndex in t.box.p1.indices { - // let _v = t.box.p1[laneIndex] - // if _v < iPosition[laneIndex] - deltaR /* True if no overlap */ { - // return false - // } - // } - let p0Flag = t.box.p0 .> (iPosition + deltaR) let p1Flag = t.box.p1 .< (iPosition - deltaR) let flag = p0Flag .| p1Flag @@ -161,19 +168,17 @@ extension Kinetics { if flag[laneIndex] { return false } - // let _v = t.box.p1[laneIndex] - // if (t.box.p0[laneIndex] > iPosition[laneIndex] + deltaR) - // || (t.box.p1[laneIndex] < iPosition[laneIndex] - // - deltaR) /* True if no overlap */ - // { - // return false - // } } return true } } } } + + @inlinable + public func dispose() { + self.tree.deinitialize(count: 1) + } } } diff --git a/Sources/ForceSimulation/Forces/ManyBodyForce.swift b/Sources/ForceSimulation/Forces/ManyBodyForce.swift index f5e6455..76122fc 100644 --- a/Sources/ForceSimulation/Forces/ManyBodyForce.swift +++ b/Sources/ForceSimulation/Forces/ManyBodyForce.swift @@ -1,3 +1,5 @@ +import simd + public struct MassCentroidKDTreeDelegate: KDTreeDelegate where Vector: SimulatableVector { @@ -82,8 +84,6 @@ extension Kinetics { public var mass: NodeMass @usableFromInline var precalculatedMass: UnsafeArray! = nil - // @usableFromInline var forces: [Vector] = [] - @inlinable public init( strength: Vector.Scalar, @@ -94,11 +94,13 @@ extension Kinetics { self.mass = nodeMass self.theta = theta self.theta2 = theta * theta + } @inlinable public func apply() { - + + // Avoid capturing self let alpha = self.kinetics.alpha let theta2 = self.theta2 let distanceMin2 = self.distanceMin2 @@ -107,16 +109,18 @@ extension Kinetics { let precalculatedMass = self.precalculatedMass.mutablePointer let positionBufferPointer = kinetics.position.mutablePointer let random = kinetics.randomGenerator + let tree = self.tree! - var tree = KDTree( - covering: self.kinetics.position, - rootDelegate: MassCentroidKDTreeDelegate(massProvider: precalculatedMass) - ) + let coveringBox = KDBox.cover(of: self.kinetics.position) + tree.pointee.reset(rootBox: coveringBox, rootDelegate: .init(massProvider: precalculatedMass)) + for p in kinetics.range { + tree.pointee.add(nodeIndex: p, at: positionBufferPointer[p]) + } for i in self.kinetics.range { let pos = positionBufferPointer[i] var f = Vector.zero - tree.visit { t in + tree.pointee.visit { t in guard t.delegate.accumulatedCount > 0 else { return false } let centroid = @@ -148,13 +152,13 @@ extension Kinetics { f += vec * k return false - } else if t.children != nil { + } else if t.childrenBufferPointer != nil { return true } if t.isFilledLeaf { - if t.nodeIndices.contains(i) { return false } + if t.nodeIndices!.contains(i) { return false } let massAcc = t.delegate.accumulatedMass @@ -176,8 +180,29 @@ extension Kinetics { public mutating func bindKinetics(_ kinetics: Kinetics) { self.kinetics = kinetics self.precalculatedMass = self.mass.calculateUnsafe(for: (kinetics.validCount)) - // self.forces = .init(repeating: .zero, count: kinetics.validCount) + self.tree = .allocate(capacity: 1) + self.tree.initialize( + to: + BufferedKDTree( + rootBox: .init( + p0: .init(repeating: 0), + p1: .init(repeating: 1) + ), + nodeCapacity: kinetics.validCount, + rootDelegate: MassCentroidKDTreeDelegate( + massProvider: precalculatedMass.mutablePointer) + ) + ) + } + + @usableFromInline + internal var tree: + UnsafeMutablePointer>>! = nil + + @inlinable + public func dispose() { + self.tree.deinitialize(count: 1) } } } diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index c5ca407..8ea6f7e 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -36,6 +36,16 @@ where next.deallocate() } } + + @inlinable + internal func contains(_ nodeIndex: Int) -> Bool { + if index == nodeIndex { return true } + if let next { + return next.pointee.contains(nodeIndex) + } else { + return false + } + } } public var nodeIndices: NodeIndex? public var box: KDBox @@ -59,7 +69,7 @@ where } -public struct BufferedKDTree: ~Copyable +public struct BufferedKDTree where Vector: SimulatableVector & L2NormCalculatable, Delegate: KDTreeDelegate @@ -315,7 +325,7 @@ where } @inlinable - static var directionCount: Int { 1 << Vector.scalarCount } + static internal var directionCount: Int { 1 << Vector.scalarCount } @inlinable internal func deinitializeBuffer() { @@ -328,7 +338,7 @@ where /// /// **Complexity**: `O(n*(2^n))`, where `n` is the dimension of the vector. @inlinable - func getIndexInChildren(_ point: Vector, relativeTo originalPoint: Vector) -> Int { + internal func getIndexInChildren(_ point: Vector, relativeTo originalPoint: Vector) -> Int { var index = 0 let mask = point .>= originalPoint @@ -660,6 +670,56 @@ extension KDTree { } + + +extension BufferedKDTree { + + /// The bounding box of the current node + @inlinable public var extent: Box { self.root.box } + + /// Visit the tree in pre-order. + /// + /// - Parameter shouldVisitChildren: a closure that returns a boolean value indicating whether should continue to visit children. + @inlinable public func visit( + shouldVisitChildren: (inout KDTreeNode) -> Bool + ) { + rootPointer.pointee.visit(shouldVisitChildren: shouldVisitChildren) + } + +} + +extension KDTreeNode { + /// Returns true is the current tree node is leaf. + /// + /// Does not guarantee that the tree node has point in it. + @inlinable public var isLeaf: Bool { childrenBufferPointer == nil } + + /// Returns true is the current tree node is internal. + /// + /// Internal tree node are always empty and do not contain any points. + @inlinable public var isInternalNode: Bool { childrenBufferPointer != nil } + + /// Returns true is the current tree node is leaf and has point in it. + @inlinable public var isFilledLeaf: Bool { nodeIndices != nil } + + /// Returns true is the current tree node is leaf and does not have point in it. + @inlinable public var isEmptyLeaf: Bool { nodeIndices == nil } + + /// Visit the tree in pre-order. + /// + /// - Parameter shouldVisitChildren: a closure that returns a boolean value indicating whether should continue to visit children. + @inlinable public mutating func visit( + shouldVisitChildren: (inout KDTreeNode) -> Bool + ) { + if shouldVisitChildren(&self) && childrenBufferPointer != nil { + // this is an internal node + for i in 0...directionCount { + childrenBufferPointer![i].visit(shouldVisitChildren: shouldVisitChildren) + } + } + } +} + // public struct KDTreeRoot // where // Vector: SimulatableVector & L2NormCalculatable, diff --git a/Sources/ForceSimulation/Simulation.swift b/Sources/ForceSimulation/Simulation.swift index 2331fb5..a3e25bb 100644 --- a/Sources/ForceSimulation/Simulation.swift +++ b/Sources/ForceSimulation/Simulation.swift @@ -1,7 +1,7 @@ /// An any-dimensional force simulation. /// The points are placed in a space where you use a SIMD data structure /// to describe their coordinates. -public struct Simulation +public final class Simulation where Vector: SimulatableVector & L2NormCalculatable, ForceField: ForceProtocol { @usableFromInline @@ -55,6 +55,11 @@ where Vector: SimulatableVector & L2NormCalculatable, ForceField: ForceProtocol< self.kinetics.updatePositions() } } + + + deinit { + self.forceField.dispose() + } } diff --git a/Tests/ForceSimulationTests/GKTreeCompareTest.swift b/Tests/ForceSimulationTests/GKTreeCompareTest.swift index 25637b7..d020c18 100644 --- a/Tests/ForceSimulationTests/GKTreeCompareTest.swift +++ b/Tests/ForceSimulationTests/GKTreeCompareTest.swift @@ -91,13 +91,14 @@ final class ManyBodyForceTests: XCTestCase { } measure { - var kdtree = KDTree, DummyQuadtreeDelegate>( - box: .init([-100.0, -100.0], [100.0, 100.0]), - spawnedDelegateBeingConsumed: DummyQuadtreeDelegate() + var kdtree = BufferedKDTree, DummyQuadtreeDelegate>( + rootBox: .init([-100.0, -100.0], [100.0, 100.0]), + nodeCapacity: nodes.count, + rootDelegate: DummyQuadtreeDelegate() ) for (i, node) in nodes.enumerated() { - kdtree.add(i, at: node) + kdtree.add(nodeIndex: i, at: node) } // traverse the tree diff --git a/Tests/ForceSimulationTests/MiserableGraphTest.swift b/Tests/ForceSimulationTests/MiserableGraphTest.swift index f2720cf..5fbdf59 100644 --- a/Tests/ForceSimulationTests/MiserableGraphTest.swift +++ b/Tests/ForceSimulationTests/MiserableGraphTest.swift @@ -91,7 +91,6 @@ struct MySealedForce: ForceField2D { } } - struct MyLatticeForce: ForceField2D { var force = SealedForce2D { Kinetics2D.LinkForce( @@ -118,7 +117,7 @@ struct MyForceField3D: ForceField3D { final class MiserableGraphTest: XCTestCase { func testLattice() { - + let myForce = SealedForce2D { Kinetics2D.ManyBodyForce(strength: -30) Kinetics2D.LinkForce( @@ -129,7 +128,7 @@ final class MiserableGraphTest: XCTestCase { Kinetics2D.CollideForce(radius: .constant(3)) } - let width = 30 + let width = 20 var edge = [(Int, Int)]() for i in 0.. Date: Fri, 1 Dec 2023 23:13:17 -0500 Subject: [PATCH 6/8] Add BufferedKDTree.reset() --- .../ForceSimulation/Forces/CollideForce.swift | 10 +++- Sources/ForceSimulation/KDTree/KDBox.swift | 6 ++ Sources/ForceSimulation/KDTree/KDTree.swift | 59 ++++++++++--------- .../MiserableGraphTest.swift | 4 +- 4 files changed, 48 insertions(+), 31 deletions(-) diff --git a/Sources/ForceSimulation/Forces/CollideForce.swift b/Sources/ForceSimulation/Forces/CollideForce.swift index 31600b4..ead0e62 100644 --- a/Sources/ForceSimulation/Forces/CollideForce.swift +++ b/Sources/ForceSimulation/Forces/CollideForce.swift @@ -104,14 +104,22 @@ extension Kinetics { for _ in 0...cover(of: self.kinetics.position) + tree.pointee.reset( rootBox: coveringBox, rootDelegate: .init(radiusBufferPointer: calculatedRadius) ) assert(tree.pointee.validCount == 1) + + for p in kinetics.range { -// print("==\(tree.pointee.validCount)") +// #if DEBUG +// let validCountBeforeAdd = tree.pointee.validCount +// #endif tree.pointee.add(nodeIndex: p, at: positionBufferPointer[p]) +// #if DEBUG +// assert(validCountBeforeAdd >= tree.pointee.validCount - 8) +// #endif } for i in kinetics.range { diff --git a/Sources/ForceSimulation/KDTree/KDBox.swift b/Sources/ForceSimulation/KDTree/KDBox.swift index a9cb0dd..8b0bace 100644 --- a/Sources/ForceSimulation/KDTree/KDBox.swift +++ b/Sources/ForceSimulation/KDTree/KDBox.swift @@ -193,6 +193,12 @@ extension KDBox { // } // } } + #if DEBUG + let testBox = Self(_p0, _p1) + for i in points.range { + assert(testBox.contains(points[i])) + } + #endif return Self(_p0, _p1) } diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index 8ea6f7e..da683a6 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -138,35 +138,30 @@ where } @inlinable - public mutating func discardAndResize( - rootBox: Box, - rootDelegate: @autoclosure () -> Delegate, - forNodeCapacity capacity: Int + internal mutating func resize( + to newTreeNodeBufferSize: Int ) { - let maxBufferCount = (capacity << Vector.scalarCount) + 1 - - let zeroNode: TreeNode = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: rootDelegate(), - box: rootBox - ) - - // should be deinitialized here (ManagedBuffer) - treeNodeBuffer = .createBuffer( - withHeader: maxBufferCount, - count: maxBufferCount, - initialValue: zeroNode + assert(newTreeNodeBufferSize >= treeNodeBuffer.header) + let newTreeNodeBuffer = UnsafeArray.createBuffer( + withHeader: newTreeNodeBufferSize, + count: newTreeNodeBufferSize, + initialValue: .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: root.delegate, + box: root.box + ) ) + newTreeNodeBuffer.withUnsafeMutablePointerToElements { + $0.moveInitialize(from: treeNodeBuffer.withUnsafeMutablePointerToElements{$0}, count: validCount) + } + treeNodeBuffer = newTreeNodeBuffer rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } + } - rootPointer.pointee = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: rootDelegate(), - box: rootBox - ) - self.validCount = 1 + @inlinable + internal mutating func resize(by factor: Int) { + resize(to: treeNodeBuffer.header * factor) } @inlinable @@ -208,6 +203,10 @@ where let spawnedDelegate = treeNode.pointee.delegate.spawn() let center = treeNode.pointee.box.center + if validCount + Self.directionCount > treeNodeBuffer.header { + resize(by: 2) + } + for j in 0.. treeNodeBuffer.header { + resize(by: 2) + } + for j in 0.. Date: Fri, 1 Dec 2023 23:38:51 -0500 Subject: [PATCH 7/8] Add BufferedKDTree --- .../ForceSimulation/Forces/CollideForce.swift | 112 +++++++++++------- Sources/ForceSimulation/KDTree/KDTree.swift | 95 ++++++++------- 2 files changed, 125 insertions(+), 82 deletions(-) diff --git a/Sources/ForceSimulation/Forces/CollideForce.swift b/Sources/ForceSimulation/Forces/CollideForce.swift index ead0e62..955e206 100644 --- a/Sources/ForceSimulation/Forces/CollideForce.swift +++ b/Sources/ForceSimulation/Forces/CollideForce.swift @@ -77,7 +77,8 @@ extension Kinetics { p1: .init(repeating: 1) ), nodeCapacity: kinetics.validCount, - rootDelegate: .init(radiusBufferPointer: self.calculatedRadius.mutablePointer) + rootDelegate: .init( + radiusBufferPointer: self.calculatedRadius.mutablePointer) ) ) } @@ -104,22 +105,21 @@ extension Kinetics { for _ in 0...cover(of: self.kinetics.position) - + tree.pointee.reset( - rootBox: coveringBox, + rootBox: coveringBox, rootDelegate: .init(radiusBufferPointer: calculatedRadius) ) assert(tree.pointee.validCount == 1) - - + for p in kinetics.range { -// #if DEBUG -// let validCountBeforeAdd = tree.pointee.validCount -// #endif + // #if DEBUG + // let validCountBeforeAdd = tree.pointee.validCount + // #endif tree.pointee.add(nodeIndex: p, at: positionBufferPointer[p]) -// #if DEBUG -// assert(validCountBeforeAdd >= tree.pointee.validCount - 8) -// #endif + // #if DEBUG + // assert(validCountBeforeAdd >= tree.pointee.validCount - 8) + // #endif } for i in kinetics.range { @@ -135,39 +135,62 @@ extension Kinetics { let maxRadiusOfQuad = t.delegate.maxNodeRadius let deltaR = maxRadiusOfQuad + iR - if let jLinkNode = t.nodeIndices { - let j = jLinkNode.index - // print("\(i)<=>\(j)") - // is leaf, make sure every collision happens once. - guard j > i else { return false } - - let jR = calculatedRadius[j] - let jOriginalPosition = positionBufferPointer[j] - let jOriginalVelocity = velocityBufferPointer[j] - var deltaPosition = - iPosition - (jOriginalPosition + jOriginalVelocity) - let l = (deltaPosition).lengthSquared() - - let deltaR = iR + jR - if l < deltaR * deltaR { - - var l = /*simd_length*/ (deltaPosition.jiggled(by: random)) - .length() - l = (deltaR - l) / l * strength - - let jR2 = jR * jR - - let k = jR2 / (iR2 + jR2) - - deltaPosition *= l - - velocityBufferPointer[i] += deltaPosition * k - velocityBufferPointer[j] -= deltaPosition * (1 - k) + if var jNode = t.nodeIndices { + while true { + let j = jNode.index + // print("\(i)<=>\(j)") + // is leaf, make sure every collision happens once. + if j > i { + + let jR = calculatedRadius[j] + let jOriginalPosition = positionBufferPointer[j] + let jOriginalVelocity = velocityBufferPointer[j] + var deltaPosition = + iPosition - (jOriginalPosition + jOriginalVelocity) + let l = (deltaPosition).lengthSquared() + + let deltaR = iR + jR + if l < deltaR * deltaR { + + var l = /*simd_length*/ (deltaPosition.jiggled(by: random)) + .length() + l = (deltaR - l) / l * strength + + let jR2 = jR * jR + + let k = jR2 / (iR2 + jR2) + + deltaPosition *= l + + velocityBufferPointer[i] += deltaPosition * k + velocityBufferPointer[j] -= deltaPosition * (1 - k) + } + } + if jNode.next == nil { + break + } else { + jNode = jNode.next!.pointee + } } - return false } + // TODO: SIMD mask + + // for laneIndex in t.box.p0.indices { + // let _v = t.box.p0[laneIndex] + // if _v > iPosition[laneIndex] + deltaR /* True if no overlap */ { + // return false + // } + // } + + // for laneIndex in t.box.p1.indices { + // let _v = t.box.p1[laneIndex] + // if _v < iPosition[laneIndex] - deltaR /* True if no overlap */ { + // return false + // } + // } + let p0Flag = t.box.p0 .> (iPosition + deltaR) let p1Flag = t.box.p1 .< (iPosition - deltaR) let flag = p0Flag .| p1Flag @@ -176,13 +199,20 @@ extension Kinetics { if flag[laneIndex] { return false } + // let _v = t.box.p1[laneIndex] + // if (t.box.p0[laneIndex] > iPosition[laneIndex] + deltaR) + // || (t.box.p1[laneIndex] < iPosition[laneIndex] + // - deltaR) /* True if no overlap */ + // { + // return false + // } } return true } } } } - + @inlinable public func dispose() { self.tree.deinitialize(count: 1) diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index f1e8b29..203586d 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -11,41 +11,6 @@ where @usableFromInline var next: UnsafeMutablePointer? - @inlinable - internal init( - nodeIndex: Int - ) { - self.index = nodeIndex - self.next = nil - } - - @inlinable - internal mutating func append(nodeIndex: Int) { - if let next { - next.pointee.append(nodeIndex: nodeIndex) - } else { - next = .allocate(capacity: 1) - next!.pointee = .init(nodeIndex: nodeIndex) - } - } - - @inlinable - internal func deinitialize() { - if let next { - next.pointee.deinitialize() - next.deallocate() - } - } - - @inlinable - internal func contains(_ nodeIndex: Int) -> Bool { - if index == nodeIndex { return true } - if let next { - return next.pointee.contains(nodeIndex) - } else { - return false - } - } } public var nodeIndices: NodeIndex? public var box: KDBox @@ -78,7 +43,7 @@ where public typealias TreeNode = KDTreeNode @usableFromInline - var rootPointer: UnsafeMutablePointer + internal var rootPointer: UnsafeMutablePointer @usableFromInline internal var validCount: Int = 0 @@ -87,7 +52,7 @@ where internal var treeNodeBuffer: UnsafeArray @inlinable - var clusterDistanceSquared: Vector.Scalar { + static internal var clusterDistanceSquared: Vector.Scalar { return Vector.clusterDistanceSquared } @@ -153,7 +118,8 @@ where ) ) newTreeNodeBuffer.withUnsafeMutablePointerToElements { - $0.moveInitialize(from: treeNodeBuffer.withUnsafeMutablePointerToElements{$0}, count: validCount) + $0.moveInitialize( + from: treeNodeBuffer.withUnsafeMutablePointerToElements { $0 }, count: validCount) } treeNodeBuffer = newTreeNodeBuffer rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } @@ -193,8 +159,8 @@ where treeNode.pointee.nodeIndices = .init(nodeIndex: nodeIndex) treeNode.pointee.nodePosition = point return - } else if treeNode.pointee.nodePosition == point - || treeNode.pointee.nodePosition.distanceSquared(to: point) < clusterDistanceSquared + } else if treeNode.pointee.nodePosition.distanceSquared(to: point) + < Self.clusterDistanceSquared { treeNode.pointee.nodeIndices!.append(nodeIndex: nodeIndex) return @@ -293,7 +259,6 @@ where let newChildrenPointer = self.rootPointer + validCount - if validCount + Self.directionCount > treeNodeBuffer.header { resize(by: 2) } @@ -691,6 +656,54 @@ extension BufferedKDTree { } +extension KDTreeNode.NodeIndex { + + @inlinable + internal init( + nodeIndex: Int + ) { + self.index = nodeIndex + self.next = nil + } + + @inlinable + internal mutating func append(nodeIndex: Int) { + if let next { + next.pointee.append(nodeIndex: nodeIndex) + } else { + next = .allocate(capacity: 1) + next!.initialize(to: .init(nodeIndex: nodeIndex)) + // next!.pointee = .init(nodeIndex: nodeIndex) + } + } + + @inlinable + internal func deinitialize() { + if let next { + next.pointee.deinitialize() + next.deallocate() + } + } + + @inlinable + internal func contains(_ nodeIndex: Int) -> Bool { + if index == nodeIndex { return true } + if let next { + return next.pointee.contains(nodeIndex) + } else { + return false + } + } + + @inlinable + internal func forEach(_ body: (Int) -> Void) { + body(index) + if let next { + next.pointee.forEach(body) + } + } +} + extension KDTreeNode { /// Returns true is the current tree node is leaf. /// From 3f2d6b16d01b207289e238b6f3a34ee55aa47c20 Mon Sep 17 00:00:00 2001 From: "Li, Zhen" Date: Fri, 1 Dec 2023 23:43:47 -0500 Subject: [PATCH 8/8] Refactor --- README.md | 4 +- .../KDTree/BufferedKDTree.swift | 419 +++++++++++++++++ Sources/ForceSimulation/KDTree/KDTree.swift | 443 ------------------ 3 files changed, 421 insertions(+), 445 deletions(-) create mode 100644 Sources/ForceSimulation/KDTree/BufferedKDTree.swift diff --git a/README.md b/README.md index 9493d95..38c4946 100644 --- a/README.md +++ b/README.md @@ -210,9 +210,9 @@ See [Example](https://github.com/li3zhen1/Grape/tree/main/Examples/ForceDirected #### Simulation -Grape uses simd to calculate position and velocity. Currently it takes **~0.010** seconds to iterate 120 times over the example graph(2D). (77 vertices, 254 edges, with manybody, center, collide and link forces. Release build on a M1 Max, tested with command `swift test -c release`) +Grape uses simd to calculate position and velocity. Currently it takes **~0.005** seconds to iterate 120 times over the example graph(2D). (77 vertices, 254 edges, with manybody, center, collide and link forces. Release build on a M1 Max, tested with command `swift test -c release`) -For 3D simulation, it takes **~0.014** seconds for the same graph and same configs. +For 3D simulation, it takes **~0.008** seconds for the same graph and same configs. > [!IMPORTANT] > Due to heavy use of generics (some are not specialized in Debug mode), the performance in Debug build is ~100x slower than Release build. Grape might ship a version with pre-inlined generics to address this problem. diff --git a/Sources/ForceSimulation/KDTree/BufferedKDTree.swift b/Sources/ForceSimulation/KDTree/BufferedKDTree.swift new file mode 100644 index 0000000..0d6ef4e --- /dev/null +++ b/Sources/ForceSimulation/KDTree/BufferedKDTree.swift @@ -0,0 +1,419 @@ +public struct KDTreeNode +where + Vector: SimulatableVector & L2NormCalculatable, + Delegate: KDTreeDelegate +{ + public struct NodeIndex { + + @usableFromInline + var index: Int + + @usableFromInline + var next: UnsafeMutablePointer? + + } + public var nodeIndices: NodeIndex? + public var box: KDBox + public var nodePosition: Vector + public var childrenBufferPointer: UnsafeMutablePointer>? + public var delegate: Delegate + + @inlinable + init( + nodeIndices: NodeIndex?, + childrenBufferPointer: UnsafeMutablePointer>?, + delegate: consuming Delegate, + box: consuming KDBox + ) { + self.childrenBufferPointer = childrenBufferPointer + self.nodeIndices = nodeIndices + self.delegate = consume delegate + self.box = consume box + self.nodePosition = .zero + } + +} + +public struct BufferedKDTree +where + Vector: SimulatableVector & L2NormCalculatable, + Delegate: KDTreeDelegate +{ + public typealias Box = KDBox + public typealias TreeNode = KDTreeNode + + @usableFromInline + internal var rootPointer: UnsafeMutablePointer + + @usableFromInline + internal var validCount: Int = 0 + + @usableFromInline + internal var treeNodeBuffer: UnsafeArray + + @inlinable + static internal var clusterDistanceSquared: Vector.Scalar { + return Vector.clusterDistanceSquared + } + + @inlinable + public var root: TreeNode { rootPointer.pointee } + + @inlinable + public init( + rootBox: Box, + nodeCapacity: Int, + rootDelegate: @autoclosure () -> Delegate + ) { + let maxBufferCount = (nodeCapacity << Vector.scalarCount) + 1 + let zeroNode: TreeNode = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate(), + box: rootBox + ) + treeNodeBuffer = .createBuffer( + withHeader: maxBufferCount, + count: maxBufferCount, + initialValue: zeroNode + ) + rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } + + rootPointer.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate(), + box: rootBox + ) + self.validCount = 1 + } + + @inlinable + public mutating func reset( + rootBox: Box, + rootDelegate: @autoclosure () -> Delegate + ) { + rootPointer.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: rootDelegate(), + box: rootBox + ) + self.validCount = 1 + } + + @inlinable + internal mutating func resize( + to newTreeNodeBufferSize: Int + ) { + assert(newTreeNodeBufferSize >= treeNodeBuffer.header) + let newTreeNodeBuffer = UnsafeArray.createBuffer( + withHeader: newTreeNodeBufferSize, + count: newTreeNodeBufferSize, + initialValue: .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: root.delegate, + box: root.box + ) + ) + newTreeNodeBuffer.withUnsafeMutablePointerToElements { + $0.moveInitialize( + from: treeNodeBuffer.withUnsafeMutablePointerToElements { $0 }, count: validCount) + } + treeNodeBuffer = newTreeNodeBuffer + rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } + } + + @inlinable + internal mutating func resize(by factor: Int) { + resize(to: treeNodeBuffer.header * factor) + } + + @inlinable + public mutating func add( + nodeIndex: Int, + at point: Vector + ) { + assert(validCount > 0) + cover(point: point) + addWithoutCover( + onTreeNode: rootPointer, + nodeOf: nodeIndex, + at: point + ) + } + + @inlinable + internal mutating func addWithoutCover( + onTreeNode treeNode: UnsafeMutablePointer, + nodeOf nodeIndex: Int, + at point: Vector + ) { + + defer { + treeNode.pointee.delegate.didAddNode(nodeIndex, at: point) + } + guard treeNode.pointee.childrenBufferPointer != nil else { + if treeNode.pointee.nodeIndices == nil { + treeNode.pointee.nodeIndices = .init(nodeIndex: nodeIndex) + treeNode.pointee.nodePosition = point + return + } else if treeNode.pointee.nodePosition.distanceSquared(to: point) + < Self.clusterDistanceSquared + { + treeNode.pointee.nodeIndices!.append(nodeIndex: nodeIndex) + return + } else { + + let spawnedDelegate = treeNode.pointee.delegate.spawn() + let center = treeNode.pointee.box.center + + if validCount + Self.directionCount > treeNodeBuffer.header { + resize(by: 2) + } + + for j in 0..> i) & 0b1 + // TODO: use simd mask + if isOnTheHigherRange != 0 { + __box.p0[i] = center[i] + } else { + __box.p1[i] = center[i] + } + } + + treeNodeBuffer[validCount + j] = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: spawnedDelegate, + box: __box + ) + + } + treeNode.pointee.childrenBufferPointer = rootPointer + validCount + validCount += Self.directionCount + + if let childrenBufferPointer = treeNode.pointee.childrenBufferPointer { + let direction = getIndexInChildren( + treeNode.pointee.nodePosition, + relativeTo: center + ) + + childrenBufferPointer[direction].nodeIndices = treeNode.pointee.nodeIndices + childrenBufferPointer[direction].nodePosition = treeNode.pointee.nodePosition + childrenBufferPointer[direction].delegate = treeNode.pointee.delegate + treeNode.pointee.nodeIndices = nil + treeNode.pointee.nodePosition = .zero + } + + let directionOfNewNode = getIndexInChildren(point, relativeTo: center) + // spawnedChildren[directionOfNewNode].addWithoutCover(nodeIndex, at: point) + addWithoutCover( + onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, + nodeOf: nodeIndex, + at: point + ) + + // self.children = spawnedChildren + return + + } + } + + let directionOfNewNode = getIndexInChildren(point, relativeTo: treeNode.pointee.box.center) + self.addWithoutCover( + onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, + nodeOf: nodeIndex, + at: point + ) + return + } + + @inlinable + internal mutating func cover(point: Vector) { + if self.root.box.contains(point) { return } + + repeat { + let direction = self.getIndexInChildren(point, relativeTo: self.root.box.p0) + self.expand(towards: direction) + } while !self.root.box.contains(point) + } + + @inlinable + internal mutating func expand(towards direction: Int) { + let nailedDirection = (Self.directionCount - 1) - direction + let nailedCorner = self.root.box.getCorner(of: nailedDirection) + let _corner = self.root.box.getCorner(of: direction) + let expandedCorner = (_corner + _corner) - nailedCorner + let newRootBox = Box(nailedCorner, expandedCorner) + + let _rootValue = self.root + + // spawn the delegate with the same internal values + // for the children, use implicit copy of spawned + let spawned = _rootValue.delegate.spawn() + + let newChildrenPointer = self.rootPointer + validCount + + if validCount + Self.directionCount > treeNodeBuffer.header { + resize(by: 2) + } + + for j in 0..> i) & 0b1 + // TODO: use simd mask + if isOnTheHigherRange != 0 { + __box.p0[i] = _corner[i] + } else { + __box.p1[i] = _corner[i] + } + } + + self.treeNodeBuffer[validCount + j] = .init( + nodeIndices: nil, + childrenBufferPointer: nil, + delegate: j != nailedDirection ? _rootValue.delegate : spawned, + box: __box + ) + } + self.validCount += Self.directionCount + + self.rootPointer.pointee = .init( + nodeIndices: nil, + childrenBufferPointer: newChildrenPointer, + delegate: _rootValue.delegate, + box: newRootBox + ) + } + + @inlinable + static internal var directionCount: Int { 1 << Vector.scalarCount } + + @inlinable + internal func deinitializeBuffer() { + _ = treeNodeBuffer.withUnsafeMutablePointerToElements { + $0.deinitialize(count: Self.directionCount) + } + } + + /// Get the index of the child that contains the point. + /// + /// **Complexity**: `O(n*(2^n))`, where `n` is the dimension of the vector. + @inlinable + internal func getIndexInChildren(_ point: Vector, relativeTo originalPoint: Vector) -> Int { + var index = 0 + + let mask = point .>= originalPoint + + for i in 0..) -> Bool + ) { + rootPointer.pointee.visit(shouldVisitChildren: shouldVisitChildren) + } + +} + +extension KDTreeNode.NodeIndex { + + @inlinable + internal init( + nodeIndex: Int + ) { + self.index = nodeIndex + self.next = nil + } + + @inlinable + internal mutating func append(nodeIndex: Int) { + if let next { + next.pointee.append(nodeIndex: nodeIndex) + } else { + next = .allocate(capacity: 1) + next!.initialize(to: .init(nodeIndex: nodeIndex)) + // next!.pointee = .init(nodeIndex: nodeIndex) + } + } + + @inlinable + internal func deinitialize() { + if let next { + next.pointee.deinitialize() + next.deallocate() + } + } + + @inlinable + internal func contains(_ nodeIndex: Int) -> Bool { + if index == nodeIndex { return true } + if let next { + return next.pointee.contains(nodeIndex) + } else { + return false + } + } + + @inlinable + internal func forEach(_ body: (Int) -> Void) { + body(index) + if let next { + next.pointee.forEach(body) + } + } +} + +extension KDTreeNode { + /// Returns true is the current tree node is leaf. + /// + /// Does not guarantee that the tree node has point in it. + @inlinable public var isLeaf: Bool { childrenBufferPointer == nil } + + /// Returns true is the current tree node is internal. + /// + /// Internal tree node are always empty and do not contain any points. + @inlinable public var isInternalNode: Bool { childrenBufferPointer != nil } + + /// Returns true is the current tree node is leaf and has point in it. + @inlinable public var isFilledLeaf: Bool { nodeIndices != nil } + + /// Returns true is the current tree node is leaf and does not have point in it. + @inlinable public var isEmptyLeaf: Bool { nodeIndices == nil } + + /// Visit the tree in pre-order. + /// + /// - Parameter shouldVisitChildren: a closure that returns a boolean value indicating whether should continue to visit children. + @inlinable public mutating func visit( + shouldVisitChildren: (inout KDTreeNode) -> Bool + ) { + if shouldVisitChildren(&self) && childrenBufferPointer != nil { + // this is an internal node + for i in 0...directionCount { + childrenBufferPointer![i].visit(shouldVisitChildren: shouldVisitChildren) + } + } + } +} \ No newline at end of file diff --git a/Sources/ForceSimulation/KDTree/KDTree.swift b/Sources/ForceSimulation/KDTree/KDTree.swift index 203586d..3572843 100644 --- a/Sources/ForceSimulation/KDTree/KDTree.swift +++ b/Sources/ForceSimulation/KDTree/KDTree.swift @@ -1,326 +1,3 @@ -public struct KDTreeNode -where - Vector: SimulatableVector & L2NormCalculatable, - Delegate: KDTreeDelegate -{ - public struct NodeIndex { - - @usableFromInline - var index: Int - - @usableFromInline - var next: UnsafeMutablePointer? - - } - public var nodeIndices: NodeIndex? - public var box: KDBox - public var nodePosition: Vector - public var childrenBufferPointer: UnsafeMutablePointer>? - public var delegate: Delegate - - @inlinable - init( - nodeIndices: NodeIndex?, - childrenBufferPointer: UnsafeMutablePointer>?, - delegate: consuming Delegate, - box: consuming KDBox - ) { - self.childrenBufferPointer = childrenBufferPointer - self.nodeIndices = nodeIndices - self.delegate = consume delegate - self.box = consume box - self.nodePosition = .zero - } - -} - -public struct BufferedKDTree -where - Vector: SimulatableVector & L2NormCalculatable, - Delegate: KDTreeDelegate -{ - public typealias Box = KDBox - public typealias TreeNode = KDTreeNode - - @usableFromInline - internal var rootPointer: UnsafeMutablePointer - - @usableFromInline - internal var validCount: Int = 0 - - @usableFromInline - internal var treeNodeBuffer: UnsafeArray - - @inlinable - static internal var clusterDistanceSquared: Vector.Scalar { - return Vector.clusterDistanceSquared - } - - @inlinable - public var root: TreeNode { rootPointer.pointee } - - @inlinable - public init( - rootBox: Box, - nodeCapacity: Int, - rootDelegate: @autoclosure () -> Delegate - ) { - let maxBufferCount = (nodeCapacity << Vector.scalarCount) + 1 - let zeroNode: TreeNode = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: rootDelegate(), - box: rootBox - ) - treeNodeBuffer = .createBuffer( - withHeader: maxBufferCount, - count: maxBufferCount, - initialValue: zeroNode - ) - rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } - - rootPointer.pointee = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: rootDelegate(), - box: rootBox - ) - self.validCount = 1 - } - - @inlinable - public mutating func reset( - rootBox: Box, - rootDelegate: @autoclosure () -> Delegate - ) { - rootPointer.pointee = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: rootDelegate(), - box: rootBox - ) - self.validCount = 1 - } - - @inlinable - internal mutating func resize( - to newTreeNodeBufferSize: Int - ) { - assert(newTreeNodeBufferSize >= treeNodeBuffer.header) - let newTreeNodeBuffer = UnsafeArray.createBuffer( - withHeader: newTreeNodeBufferSize, - count: newTreeNodeBufferSize, - initialValue: .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: root.delegate, - box: root.box - ) - ) - newTreeNodeBuffer.withUnsafeMutablePointerToElements { - $0.moveInitialize( - from: treeNodeBuffer.withUnsafeMutablePointerToElements { $0 }, count: validCount) - } - treeNodeBuffer = newTreeNodeBuffer - rootPointer = treeNodeBuffer.withUnsafeMutablePointerToElements { $0 } - } - - @inlinable - internal mutating func resize(by factor: Int) { - resize(to: treeNodeBuffer.header * factor) - } - - @inlinable - public mutating func add( - nodeIndex: Int, - at point: Vector - ) { - assert(validCount > 0) - cover(point: point) - addWithoutCover( - onTreeNode: rootPointer, - nodeOf: nodeIndex, - at: point - ) - } - - @inlinable - internal mutating func addWithoutCover( - onTreeNode treeNode: UnsafeMutablePointer, - nodeOf nodeIndex: Int, - at point: Vector - ) { - - defer { - treeNode.pointee.delegate.didAddNode(nodeIndex, at: point) - } - guard treeNode.pointee.childrenBufferPointer != nil else { - if treeNode.pointee.nodeIndices == nil { - treeNode.pointee.nodeIndices = .init(nodeIndex: nodeIndex) - treeNode.pointee.nodePosition = point - return - } else if treeNode.pointee.nodePosition.distanceSquared(to: point) - < Self.clusterDistanceSquared - { - treeNode.pointee.nodeIndices!.append(nodeIndex: nodeIndex) - return - } else { - - let spawnedDelegate = treeNode.pointee.delegate.spawn() - let center = treeNode.pointee.box.center - - if validCount + Self.directionCount > treeNodeBuffer.header { - resize(by: 2) - } - - for j in 0..> i) & 0b1 - // TODO: use simd mask - if isOnTheHigherRange != 0 { - __box.p0[i] = center[i] - } else { - __box.p1[i] = center[i] - } - } - - treeNodeBuffer[validCount + j] = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: spawnedDelegate, - box: __box - ) - - } - treeNode.pointee.childrenBufferPointer = rootPointer + validCount - validCount += Self.directionCount - - if let childrenBufferPointer = treeNode.pointee.childrenBufferPointer { - let direction = getIndexInChildren( - treeNode.pointee.nodePosition, - relativeTo: center - ) - - childrenBufferPointer[direction].nodeIndices = treeNode.pointee.nodeIndices - childrenBufferPointer[direction].nodePosition = treeNode.pointee.nodePosition - childrenBufferPointer[direction].delegate = treeNode.pointee.delegate - treeNode.pointee.nodeIndices = nil - treeNode.pointee.nodePosition = .zero - } - - let directionOfNewNode = getIndexInChildren(point, relativeTo: center) - // spawnedChildren[directionOfNewNode].addWithoutCover(nodeIndex, at: point) - addWithoutCover( - onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, - nodeOf: nodeIndex, - at: point - ) - - // self.children = spawnedChildren - return - - } - } - - let directionOfNewNode = getIndexInChildren(point, relativeTo: treeNode.pointee.box.center) - self.addWithoutCover( - onTreeNode: treeNode.pointee.childrenBufferPointer! + directionOfNewNode, - nodeOf: nodeIndex, - at: point - ) - return - } - - @inlinable - internal mutating func cover(point: Vector) { - if self.root.box.contains(point) { return } - - repeat { - let direction = self.getIndexInChildren(point, relativeTo: self.root.box.p0) - self.expand(towards: direction) - } while !self.root.box.contains(point) - } - - @inlinable - internal mutating func expand(towards direction: Int) { - let nailedDirection = (Self.directionCount - 1) - direction - let nailedCorner = self.root.box.getCorner(of: nailedDirection) - let _corner = self.root.box.getCorner(of: direction) - let expandedCorner = (_corner + _corner) - nailedCorner - let newRootBox = Box(nailedCorner, expandedCorner) - - let _rootValue = self.root - - // spawn the delegate with the same internal values - // for the children, use implicit copy of spawned - let spawned = _rootValue.delegate.spawn() - - let newChildrenPointer = self.rootPointer + validCount - - if validCount + Self.directionCount > treeNodeBuffer.header { - resize(by: 2) - } - - for j in 0..> i) & 0b1 - // TODO: use simd mask - if isOnTheHigherRange != 0 { - __box.p0[i] = _corner[i] - } else { - __box.p1[i] = _corner[i] - } - } - - self.treeNodeBuffer[validCount + j] = .init( - nodeIndices: nil, - childrenBufferPointer: nil, - delegate: j != nailedDirection ? _rootValue.delegate : spawned, - box: __box - ) - } - self.validCount += Self.directionCount - - self.rootPointer.pointee = .init( - nodeIndices: nil, - childrenBufferPointer: newChildrenPointer, - delegate: _rootValue.delegate, - box: newRootBox - ) - } - - @inlinable - static internal var directionCount: Int { 1 << Vector.scalarCount } - - @inlinable - internal func deinitializeBuffer() { - _ = treeNodeBuffer.withUnsafeMutablePointerToElements { - $0.deinitialize(count: Self.directionCount) - } - } - - /// Get the index of the child that contains the point. - /// - /// **Complexity**: `O(n*(2^n))`, where `n` is the dimension of the vector. - @inlinable - internal func getIndexInChildren(_ point: Vector, relativeTo originalPoint: Vector) -> Int { - var index = 0 - - let mask = point .>= originalPoint - - for i in 0..) -> Bool - ) { - rootPointer.pointee.visit(shouldVisitChildren: shouldVisitChildren) - } - -} - -extension KDTreeNode.NodeIndex { - - @inlinable - internal init( - nodeIndex: Int - ) { - self.index = nodeIndex - self.next = nil - } - - @inlinable - internal mutating func append(nodeIndex: Int) { - if let next { - next.pointee.append(nodeIndex: nodeIndex) - } else { - next = .allocate(capacity: 1) - next!.initialize(to: .init(nodeIndex: nodeIndex)) - // next!.pointee = .init(nodeIndex: nodeIndex) - } - } - - @inlinable - internal func deinitialize() { - if let next { - next.pointee.deinitialize() - next.deallocate() - } - } - - @inlinable - internal func contains(_ nodeIndex: Int) -> Bool { - if index == nodeIndex { return true } - if let next { - return next.pointee.contains(nodeIndex) - } else { - return false - } - } - - @inlinable - internal func forEach(_ body: (Int) -> Void) { - body(index) - if let next { - next.pointee.forEach(body) - } - } -} - -extension KDTreeNode { - /// Returns true is the current tree node is leaf. - /// - /// Does not guarantee that the tree node has point in it. - @inlinable public var isLeaf: Bool { childrenBufferPointer == nil } - - /// Returns true is the current tree node is internal. - /// - /// Internal tree node are always empty and do not contain any points. - @inlinable public var isInternalNode: Bool { childrenBufferPointer != nil } - - /// Returns true is the current tree node is leaf and has point in it. - @inlinable public var isFilledLeaf: Bool { nodeIndices != nil } - - /// Returns true is the current tree node is leaf and does not have point in it. - @inlinable public var isEmptyLeaf: Bool { nodeIndices == nil } - - /// Visit the tree in pre-order. - /// - /// - Parameter shouldVisitChildren: a closure that returns a boolean value indicating whether should continue to visit children. - @inlinable public mutating func visit( - shouldVisitChildren: (inout KDTreeNode) -> Bool - ) { - if shouldVisitChildren(&self) && childrenBufferPointer != nil { - // this is an internal node - for i in 0...directionCount { - childrenBufferPointer![i].visit(shouldVisitChildren: shouldVisitChildren) - } - } - } -} - -// public struct KDTreeRoot -// where -// Vector: SimulatableVector & L2NormCalculatable, -// Delegate: KDTreeDelegate -// { -// public var rootPointer: KDTree -// @usableFromInline let propertyBuffer: UnsafeMutablePointer - -// @inlinable -// public init( -// rootPointer: KDTree, -// propertyBuffer: UnsafeMutablePointer -// ) { -// self.rootPointer = rootPointer -// self.propertyBuffer = propertyBuffer -// } - -// @inlinable -// public mutating func add(_ nodeIndex: Int, at point: Vector) { -// rootPointer.cover(point) -// rootPointer.addWithoutCover(nodeIndex, at: point) -// } -// }