From 91e2c7ce39ee753a232deaf37f08a0cbb45a06ed Mon Sep 17 00:00:00 2001 From: Vishesh Yadav Date: Fri, 18 Aug 2023 18:22:49 -0700 Subject: [PATCH] Move around and some cleanup --- Sources/ARTreeModule/ARTree+get.swift | 9 +-- ...e+basics.swift => InternalNode+impl.swift} | 0 Sources/ARTreeModule/Node+Storage.swift | 10 +-- Sources/ARTreeModule/Node.swift | 60 ++---------------- Sources/ARTreeModule/Node16.swift | 55 ++++++++-------- Sources/ARTreeModule/Node256.swift | 46 +++++++------- Sources/ARTreeModule/Node4.swift | 53 ++++++++-------- Sources/ARTreeModule/Node48.swift | 62 +++++++++---------- Sources/ARTreeModule/NodeLeaf.swift | 46 ++++++++------ Sources/ARTreeModule/RawNode.swift | 53 ++++++++++++++++ 10 files changed, 194 insertions(+), 200 deletions(-) rename Sources/ARTreeModule/{Node+basics.swift => InternalNode+impl.swift} (100%) create mode 100644 Sources/ARTreeModule/RawNode.swift diff --git a/Sources/ARTreeModule/ARTree+get.swift b/Sources/ARTreeModule/ARTree+get.swift index ee9aa08da..01ec78b9f 100644 --- a/Sources/ARTreeModule/ARTree+get.swift +++ b/Sources/ARTreeModule/ARTree+get.swift @@ -11,21 +11,22 @@ extension ARTree { public func getValue(key: Key) -> Value? { + assert(root != nil, "root can't be nil") var current = root var depth = 0 while depth <= key.count { - guard let _node = current else { + guard let _rawNode = current else { return nil } - if _node.type == .leaf { - let leaf: NodeLeaf = _node.toLeafNode() + if _rawNode.type == .leaf { + let leaf: NodeLeaf = _rawNode.toLeafNode() return leaf.keyEquals(with: key) ? leaf.value() : nil } - let node = _node.toInternalNode() + let node = _rawNode.toInternalNode() if node.partialLength > 0 { let prefixLen = node.prefixMismatch(withKey: key, fromIndex: depth) assert(prefixLen <= Const.maxPartialLength, "partial length is always bounded") diff --git a/Sources/ARTreeModule/Node+basics.swift b/Sources/ARTreeModule/InternalNode+impl.swift similarity index 100% rename from Sources/ARTreeModule/Node+basics.swift rename to Sources/ARTreeModule/InternalNode+impl.swift diff --git a/Sources/ARTreeModule/Node+Storage.swift b/Sources/ARTreeModule/Node+Storage.swift index d36f22a8d..76c08b83c 100644 --- a/Sources/ARTreeModule/Node+Storage.swift +++ b/Sources/ARTreeModule/Node+Storage.swift @@ -13,7 +13,7 @@ typealias RawNodeBuffer = ManagedBuffer final class NodeBuffer: RawNodeBuffer { deinit { - Mn.deinitialize(NodeStorage(self)) + Mn.deinitialize(NodeStorage(buf: self)) } } @@ -22,16 +22,10 @@ struct NodeStorage { } extension NodeStorage { - fileprivate init(_ buf: NodeBuffer) { - self.buf = buf - } - init(raw: RawNodeBuffer) { self.buf = raw as! NodeBuffer } -} -extension NodeStorage { static func create(type: NodeType, size: Int) -> RawNodeBuffer { let buf = NodeBuffer.create(minimumCapacity: size, makingHeaderWith: {_ in type }) @@ -40,7 +34,9 @@ extension NodeStorage { } return buf } +} +extension NodeStorage { func withUnsafePointer(_ body: (UnsafeMutableRawPointer) throws -> R) rethrows -> R { return try buf.withUnsafeMutablePointerToElements { return try body(UnsafeMutableRawPointer($0)) diff --git a/Sources/ARTreeModule/Node.swift b/Sources/ARTreeModule/Node.swift index 84c0a8470..993180bb0 100644 --- a/Sources/ARTreeModule/Node.swift +++ b/Sources/ARTreeModule/Node.swift @@ -14,64 +14,17 @@ public typealias Key = [KeyPart] typealias ChildSlotPtr = UnsafeMutablePointer -struct RawNode { - var storage: RawNodeBuffer -} - -extension RawNode { - init(from: N) { - self.storage = from.storage.buf - } -} - -extension RawNode { - var type: NodeType { - @inline(__always) get { return storage.header } - } - - func toInternalNode() -> any InternalNode { - switch type { - case .node4: - return Node4(buffer: storage) - case .node16: - return Node16(buffer: storage) - case .node48: - return Node48(buffer: storage) - case .node256: - return Node256(buffer: storage) - default: - assert(false, "leaf nodes are not internal nodes") - } - } - - func toLeafNode() -> NodeLeaf { - assert(type == .leaf) - return NodeLeaf(ptr: storage) - } - - func toManagedNode() -> any ManagedNode { - switch type { - case .leaf: - return toLeafNode() - default: - return toInternalNode() - } - } -} - protocol ManagedNode: NodePrettyPrinter { - static func deinitialize(_ storage: NodeStorage) + typealias Storage = NodeStorage + + static func deinitialize(_ storage: NodeStorage) static var type: NodeType { get } - var storage: NodeStorage { get } + var storage: Storage { get } var type: NodeType { get } var rawNode: RawNode { get } } -extension ManagedNode { - var rawNode: RawNode { RawNode(from: self) } -} - protocol InternalNode: ManagedNode { typealias Index = Int typealias Header = InternalNodeHeader @@ -103,7 +56,6 @@ protocol InternalNode: ManagedNode { } extension ManagedNode { - static func deinitialize(_ storage: NodeStorage) { - // TODO - } + var rawNode: RawNode { RawNode(from: self) } + var type: NodeType { Self.type } } diff --git a/Sources/ARTreeModule/Node16.swift b/Sources/ARTreeModule/Node16.swift index 0a9232c24..7a427317e 100644 --- a/Sources/ARTreeModule/Node16.swift +++ b/Sources/ARTreeModule/Node16.swift @@ -10,36 +10,15 @@ //===----------------------------------------------------------------------===// struct Node16 { - typealias Storage = NodeStorage - var storage: Storage - - init(buffer: RawNodeBuffer) { - self.init(storage: Storage(raw: buffer)) - } - - init(storage: Storage) { - self.storage = storage - } } extension Node16 { - typealias Keys = UnsafeMutableBufferPointer - typealias Childs = UnsafeMutableBufferPointer - - func withBody(body: (Keys, Childs) throws -> R) rethrows -> R { - return try storage.withBodyPointer { bodyPtr in - let keys = UnsafeMutableBufferPointer( - start: bodyPtr.assumingMemoryBound(to: KeyPart.self), - count: Self.numKeys - ) - let childPtr = bodyPtr - .advanced(by: Self.numKeys * MemoryLayout.stride) - .assumingMemoryBound(to: RawNode?.self) - let childs = UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + static let type: NodeType = .node16 + static let numKeys: Int = 16 - return try body(keys, childs) - } + init(buffer: RawNodeBuffer) { + self.init(storage: Storage(raw: buffer)) } } @@ -97,12 +76,27 @@ extension Node16 { } } -extension Node16: InternalNode { - static let type: NodeType = .node16 - var type: NodeType { .node16 } +extension Node16 { + typealias Keys = UnsafeMutableBufferPointer + typealias Childs = UnsafeMutableBufferPointer - static let numKeys: Int = 16 + func withBody(body: (Keys, Childs) throws -> R) rethrows -> R { + return try storage.withBodyPointer { bodyPtr in + let keys = UnsafeMutableBufferPointer( + start: bodyPtr.assumingMemoryBound(to: KeyPart.self), + count: Self.numKeys + ) + let childPtr = bodyPtr + .advanced(by: Self.numKeys * MemoryLayout.stride) + .assumingMemoryBound(to: RawNode?.self) + let childs = UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + + return try body(keys, childs) + } + } +} +extension Node16: InternalNode { static var size: Int { MemoryLayout.stride + Self.numKeys * (MemoryLayout.stride + MemoryLayout.stride) @@ -219,4 +213,7 @@ extension Node16: InternalNode { } } } + + static func deinitialize(_ storage: NodeStorage) { + } } diff --git a/Sources/ARTreeModule/Node256.swift b/Sources/ARTreeModule/Node256.swift index c97e9b70e..11ed447ae 100644 --- a/Sources/ARTreeModule/Node256.swift +++ b/Sources/ARTreeModule/Node256.swift @@ -9,31 +9,16 @@ // //===----------------------------------------------------------------------===// -struct Node256: ManagedNode { - typealias Storage = NodeStorage - +struct Node256 { var storage: Storage - - init(buffer: RawNodeBuffer) { - self.init(storage: Storage(raw: buffer)) - } - - init(storage: Storage) { - self.storage = storage - } } extension Node256 { - typealias Keys = UnsafeMutableBufferPointer - typealias Childs = UnsafeMutableBufferPointer + static let type: NodeType = .node256 + static let numKeys: Int = 256 - func withBody(body: (Childs) throws -> R) rethrows -> R { - return try storage.withBodyPointer { - return try body( - UnsafeMutableBufferPointer( - start: $0.assumingMemoryBound(to: RawNode?.self), - count: 256)) - } + init(buffer: RawNodeBuffer) { + self.init(storage: Storage(raw: buffer)) } } @@ -68,13 +53,21 @@ extension Node256 { } } +extension Node256 { + typealias Keys = UnsafeMutableBufferPointer + typealias Childs = UnsafeMutableBufferPointer -extension Node256: InternalNode { - static let type: NodeType = .node256 - var type: NodeType { .node256 } - - static let numKeys: Int = 256 + func withBody(body: (Childs) throws -> R) rethrows -> R { + return try storage.withBodyPointer { + return try body( + UnsafeMutableBufferPointer( + start: $0.assumingMemoryBound(to: RawNode?.self), + count: 256)) + } + } +} +extension Node256: InternalNode { static var size: Int { MemoryLayout.stride + 256 * MemoryLayout.stride } @@ -148,4 +141,7 @@ extension Node256: InternalNode { } } } + + static func deinitialize(_ storage: NodeStorage) { + } } diff --git a/Sources/ARTreeModule/Node4.swift b/Sources/ARTreeModule/Node4.swift index 278be0be9..095c17caf 100644 --- a/Sources/ARTreeModule/Node4.swift +++ b/Sources/ARTreeModule/Node4.swift @@ -10,39 +10,15 @@ //===----------------------------------------------------------------------===// struct Node4 { - static let numKeys: Int = 4 - - typealias Storage = NodeStorage - var storage: Storage } extension Node4 { - init(buffer: RawNodeBuffer) { - self.init(storage: Storage(raw: buffer)) - } - static let type: NodeType = .node4 - var type: NodeType { .node4 } -} - -extension Node4 { - typealias Keys = UnsafeMutableBufferPointer - typealias Childs = UnsafeMutableBufferPointer - - func withBody(body: (Keys, Childs) throws -> R) rethrows -> R { - return try storage.withBodyPointer { bodyPtr in - let keys = UnsafeMutableBufferPointer( - start: bodyPtr.assumingMemoryBound(to: KeyPart.self), - count: Self.numKeys - ) - let childPtr = bodyPtr - .advanced(by: Self.numKeys * MemoryLayout.stride) - .assumingMemoryBound(to: RawNode?.self) - let childs = UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + static let numKeys: Int = 4 - return try body(keys, childs) - } + init(buffer: RawNodeBuffer) { + self.init(storage: Storage(raw: buffer)) } } @@ -75,6 +51,26 @@ extension Node4 { } } +extension Node4 { + typealias Keys = UnsafeMutableBufferPointer + typealias Childs = UnsafeMutableBufferPointer + + func withBody(body: (Keys, Childs) throws -> R) rethrows -> R { + return try storage.withBodyPointer { bodyPtr in + let keys = UnsafeMutableBufferPointer( + start: bodyPtr.assumingMemoryBound(to: KeyPart.self), + count: Self.numKeys + ) + let childPtr = bodyPtr + .advanced(by: Self.numKeys * MemoryLayout.stride) + .assumingMemoryBound(to: RawNode?.self) + let childs = UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + + return try body(keys, childs) + } + } +} + extension Node4: InternalNode { static var size: Int { MemoryLayout.stride + Self.numKeys @@ -189,4 +185,7 @@ extension Node4: InternalNode { } } } + + static func deinitialize(_ storage: NodeStorage) { + } } diff --git a/Sources/ARTreeModule/Node48.swift b/Sources/ARTreeModule/Node48.swift index baffdba09..a17bf5309 100644 --- a/Sources/ARTreeModule/Node48.swift +++ b/Sources/ARTreeModule/Node48.swift @@ -10,41 +10,15 @@ //===----------------------------------------------------------------------===// struct Node48 { - typealias Storage = NodeStorage - var storage: Storage - - init(buffer: RawNodeBuffer) { - self.init(storage: Storage(raw: buffer)) - } - - init(storage: Storage) { - self.storage = storage - } } - extension Node48 { - typealias Keys = UnsafeMutableBufferPointer - typealias Childs = UnsafeMutableBufferPointer - - func withBody(body: (Keys, Childs) throws -> R) rethrows -> R { - return try storage.withBodyPointer { bodyPtr in - let keys = UnsafeMutableBufferPointer( - start: bodyPtr.assumingMemoryBound(to: KeyPart.self), - count: 256 - ) - - // NOTE: Initializes each key pointer to point to a value > number of children, as 0 will - // refer to the first child. - // TODO: Can we initialize buffer using any stdlib method? - let childPtr = bodyPtr - .advanced(by: 256 * MemoryLayout.stride) - .assumingMemoryBound(to: RawNode?.self) - let childs = UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + static let type: NodeType = .node48 + static let numKeys: Int = 48 - return try body(keys, childs) - } + init(buffer: RawNodeBuffer) { + self.init(storage: Storage(raw: buffer)) } } @@ -107,13 +81,31 @@ extension Node48 { } } +extension Node48 { + typealias Keys = UnsafeMutableBufferPointer + typealias Childs = UnsafeMutableBufferPointer -extension Node48: InternalNode { - static let type: NodeType = .node48 - var type: NodeType { .node48 } + func withBody(body: (Keys, Childs) throws -> R) rethrows -> R { + return try storage.withBodyPointer { bodyPtr in + let keys = UnsafeMutableBufferPointer( + start: bodyPtr.assumingMemoryBound(to: KeyPart.self), + count: 256 + ) - static let numKeys: Int = 48 + // NOTE: Initializes each key pointer to point to a value > number of children, as 0 will + // refer to the first child. + // TODO: Can we initialize buffer using any stdlib method? + let childPtr = bodyPtr + .advanced(by: 256 * MemoryLayout.stride) + .assumingMemoryBound(to: RawNode?.self) + let childs = UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + return try body(keys, childs) + } + } +} + +extension Node48: InternalNode { static var size: Int { MemoryLayout.stride + 256*MemoryLayout.stride + Self.numKeys*MemoryLayout.stride @@ -246,4 +238,6 @@ extension Node48: InternalNode { } } + static func deinitialize(_ storage: NodeStorage) { + } } diff --git a/Sources/ARTreeModule/NodeLeaf.swift b/Sources/ARTreeModule/NodeLeaf.swift index 3076f5b12..6d789da86 100644 --- a/Sources/ARTreeModule/NodeLeaf.swift +++ b/Sources/ARTreeModule/NodeLeaf.swift @@ -16,16 +16,33 @@ struct NodeLeaf { } extension NodeLeaf { + static var type: NodeType { .leaf } + init(ptr: RawNodeBuffer) { self.init(storage: Storage(raw: ptr)) } } -extension NodeLeaf: ManagedNode { - typealias KeyPtr = UnsafeMutableBufferPointer +extension NodeLeaf { + static func allocate(key: Key, value: Value, of: Value.Type) -> Self { + let size = MemoryLayout.stride + key.count + MemoryLayout.stride + let buf = NodeStorage.create(type: .leaf, size: size) + var leaf = Self(ptr: buf) - static var type: NodeType { .leaf } - var type: NodeType { .leaf } + leaf.keyLength = key.count + leaf.withKeyValue { keyPtr, valuePtr in + key.withUnsafeBytes { + UnsafeMutableRawBufferPointer(keyPtr).copyBytes(from: $0) + } + valuePtr.pointee = value + } + + return leaf + } +} + +extension NodeLeaf { + typealias KeyPtr = UnsafeMutableBufferPointer func withKey(body: (KeyPtr) throws -> R) rethrows -> R { return try storage.withUnsafePointer { @@ -86,22 +103,6 @@ extension NodeLeaf: ManagedNode { } extension NodeLeaf { - static func allocate(key: Key, value: Value, of: Value.Type) -> Self { - let size = MemoryLayout.stride + key.count + MemoryLayout.stride - let buf = NodeStorage.create(type: .leaf, size: size) - var leaf = Self(ptr: buf) - - leaf.keyLength = key.count - leaf.withKeyValue { keyPtr, valuePtr in - key.withUnsafeBytes { - UnsafeMutableRawBufferPointer(keyPtr).copyBytes(from: $0) - } - valuePtr.pointee = value - } - - return leaf - } - func keyEquals(with key: Key, depth: Int = 0) -> Bool { if key.count != keyLength { return false @@ -133,3 +134,8 @@ extension NodeLeaf { } } } + +extension NodeLeaf: ManagedNode { + static func deinitialize(_ storage: NodeStorage) { + } +} diff --git a/Sources/ARTreeModule/RawNode.swift b/Sources/ARTreeModule/RawNode.swift new file mode 100644 index 000000000..1119e93d6 --- /dev/null +++ b/Sources/ARTreeModule/RawNode.swift @@ -0,0 +1,53 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Collections open source project +// +// Copyright (c) 2023 Apple Inc. and the Swift project authors +// Licensed under Apache License v2.0 with Runtime Library Exception +// +// See https://swift.org/LICENSE.txt for license information +// +//===----------------------------------------------------------------------===// + +struct RawNode { + var storage: RawNodeBuffer + + init(from: N) { + self.storage = from.storage.buf + } +} + +extension RawNode { + var type: NodeType { + @inline(__always) get { return storage.header } + } + + func toInternalNode() -> any InternalNode { + switch type { + case .node4: + return Node4(buffer: storage) + case .node16: + return Node16(buffer: storage) + case .node48: + return Node48(buffer: storage) + case .node256: + return Node256(buffer: storage) + default: + assert(false, "leaf nodes are not internal nodes") + } + } + + func toLeafNode() -> NodeLeaf { + assert(type == .leaf) + return NodeLeaf(ptr: storage) + } + + func toManagedNode() -> any ManagedNode { + switch type { + case .leaf: + return toLeafNode() + default: + return toInternalNode() + } + } +}