Skip to content

Commit

Permalink
Move around and some cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@vishesh
vishesh committed Aug 19, 2023
1 parent 439c6a1 commit 91e2c7c
Show file tree
Hide file tree
Showing 10 changed files with 194 additions and 200 deletions.
9 changes: 5 additions & 4 deletions Sources/ARTreeModule/ARTree+get.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
0 Sources/ARTreeModule/Node+basics.swift → Sources/ARTreeModule/InternalNode+impl.swift
View file
File renamed without changes.
10 changes: 3 additions & 7 deletions Sources/ARTreeModule/Node+Storage.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ typealias RawNodeBuffer = ManagedBuffer<NodeType, UInt8>

final class NodeBuffer<Mn: ManagedNode>: RawNodeBuffer {
deinit {
Mn.deinitialize(NodeStorage(self))
Mn.deinitialize(NodeStorage(buf: self))
}
}

Expand All @@ -22,16 +22,10 @@ struct NodeStorage<Mn: ManagedNode> {
}

extension NodeStorage {
fileprivate init(_ buf: NodeBuffer<Mn>) {
self.buf = buf
}

init(raw: RawNodeBuffer) {
self.buf = raw as! NodeBuffer<Mn>
}
}

extension NodeStorage {
static func create(type: NodeType, size: Int) -> RawNodeBuffer {
let buf = NodeBuffer<Mn>.create(minimumCapacity: size,
makingHeaderWith: {_ in type })
Expand All @@ -40,7 +34,9 @@ extension NodeStorage {
}
return buf
}
}

extension NodeStorage {
func withUnsafePointer<R>(_ body: (UnsafeMutableRawPointer) throws -> R) rethrows -> R {
return try buf.withUnsafeMutablePointerToElements {
return try body(UnsafeMutableRawPointer($0))
Expand Down
60 changes: 6 additions & 54 deletions Sources/ARTreeModule/Node.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,64 +14,17 @@ public typealias Key = [KeyPart]

typealias ChildSlotPtr = UnsafeMutablePointer<RawNode?>

struct RawNode {
var storage: RawNodeBuffer
}

extension RawNode {
init<N: ManagedNode>(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<N: ManagedNode>(_ storage: NodeStorage<N>)
typealias Storage = NodeStorage<Self>

static func deinitialize(_ storage: NodeStorage<Self>)
static var type: NodeType { get }

var storage: NodeStorage<Self> { 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
Expand Down Expand Up @@ -103,7 +56,6 @@ protocol InternalNode: ManagedNode {
}

extension ManagedNode {
static func deinitialize<N: ManagedNode>(_ storage: NodeStorage<N>) {
// TODO
}
var rawNode: RawNode { RawNode(from: self) }
var type: NodeType { Self.type }
}
55 changes: 26 additions & 29 deletions Sources/ARTreeModule/Node16.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,36 +10,15 @@
//===----------------------------------------------------------------------===//

struct Node16 {
typealias Storage = NodeStorage<Self>

var storage: Storage

init(buffer: RawNodeBuffer) {
self.init(storage: Storage(raw: buffer))
}

init(storage: Storage) {
self.storage = storage
}
}

extension Node16 {
typealias Keys = UnsafeMutableBufferPointer<KeyPart>
typealias Childs = UnsafeMutableBufferPointer<RawNode?>

func withBody<R>(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<KeyPart>.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))
}
}

Expand Down Expand Up @@ -97,12 +76,27 @@ extension Node16 {
}
}

extension Node16: InternalNode {
static let type: NodeType = .node16
var type: NodeType { .node16 }
extension Node16 {
typealias Keys = UnsafeMutableBufferPointer<KeyPart>
typealias Childs = UnsafeMutableBufferPointer<RawNode?>

static let numKeys: Int = 16
func withBody<R>(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<KeyPart>.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<InternalNodeHeader>.stride + Self.numKeys
* (MemoryLayout<KeyPart>.stride + MemoryLayout<RawNode?>.stride)
Expand Down Expand Up @@ -219,4 +213,7 @@ extension Node16: InternalNode {
}
}
}

static func deinitialize(_ storage: NodeStorage<Self>) {
}
}
46 changes: 21 additions & 25 deletions Sources/ARTreeModule/Node256.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,31 +9,16 @@
//
//===----------------------------------------------------------------------===//

struct Node256: ManagedNode {
typealias Storage = NodeStorage<Self>

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<KeyPart>
typealias Childs = UnsafeMutableBufferPointer<RawNode?>
static let type: NodeType = .node256
static let numKeys: Int = 256

func withBody<R>(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))
}
}

Expand Down Expand Up @@ -68,13 +53,21 @@ extension Node256 {
}
}

extension Node256 {
typealias Keys = UnsafeMutableBufferPointer<KeyPart>
typealias Childs = UnsafeMutableBufferPointer<RawNode?>

extension Node256: InternalNode {
static let type: NodeType = .node256
var type: NodeType { .node256 }

static let numKeys: Int = 256
func withBody<R>(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<InternalNodeHeader>.stride + 256 * MemoryLayout<RawNode?>.stride
}
Expand Down Expand Up @@ -148,4 +141,7 @@ extension Node256: InternalNode {
}
}
}

static func deinitialize(_ storage: NodeStorage<Self>) {
}
}
53 changes: 26 additions & 27 deletions Sources/ARTreeModule/Node4.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,39 +10,15 @@
//===----------------------------------------------------------------------===//

struct Node4 {
static let numKeys: Int = 4

typealias Storage = NodeStorage<Self>

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<KeyPart>
typealias Childs = UnsafeMutableBufferPointer<RawNode?>

func withBody<R>(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<KeyPart>.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))
}
}

Expand Down Expand Up @@ -75,6 +51,26 @@ extension Node4 {
}
}

extension Node4 {
typealias Keys = UnsafeMutableBufferPointer<KeyPart>
typealias Childs = UnsafeMutableBufferPointer<RawNode?>

func withBody<R>(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<KeyPart>.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<InternalNodeHeader>.stride + Self.numKeys
Expand Down Expand Up @@ -189,4 +185,7 @@ extension Node4: InternalNode {
}
}
}

static func deinitialize(_ storage: NodeStorage<Self>) {
}
}
Loading

0 comments on commit 91e2c7c

Please sign in to comment.