Skip to content

Commit 761950d

Browse files
vihanbvihanb
committed
Implement and use PackedOffsetList for UnsafePath
Using an array was very heavy, now using a stack-allocated 64-byte wide type which mimics an array.
1 parent 817c97a commit 761950d

File tree

4 files changed

+119
-72
lines changed

4 files changed

+119
-72
lines changed

‎Sources/SortedCollections/BTree/_BTree+BidirectionalCollection.swift‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ extension _BTree: BidirectionalCollection {
2626
@inlinable
2727
internal var startIndex: Index {
2828
if count == 0 { return Index(nil, forTree: self) }
29-
var depth = 0
29+
var depth: UInt8 = 0
3030
var currentNode = self.root
3131
while !currentNode.read({ $0.isLeaf }) {
3232
// TODO: figure out how to avoid the swift retain here
@@ -37,7 +37,7 @@ extension _BTree: BidirectionalCollection {
3737
let path = UnsafePath(
3838
node: currentNode,
3939
slot: 0,
40-
childSlots: Array<Int>(repeating: 0, count: depth),
40+
childSlots: PackedOffsetList(depth: depth),
4141
offset: 0
4242
)
4343

@@ -238,8 +238,7 @@ extension _BTree {
238238
/// - Returns: If found, returns a path to the element. Otherwise, `nil`.
239239
@inlinable
240240
internal func anyIndex(forKey key: Key) -> Index? {
241-
var childSlots = [Int]()
242-
childSlots.reserveCapacity(BTREE_MAX_DEPTH)
241+
var childSlots = PackedOffsetList()
243242
var node: Node? = self.root
244243

245244
while let currentNode = node {
@@ -251,7 +250,7 @@ extension _BTree {
251250
if handle.isLeaf {
252251
node = nil
253252
} else {
254-
childSlots.append(keySlot)
253+
childSlots.append(UInt16(keySlot))
255254
node = handle[childAt: keySlot]
256255
}
257256

@@ -279,8 +278,7 @@ extension _BTree {
279278
return Index(nil, forTree: self)
280279
}
281280

282-
var childSlots = [Int]()
283-
childSlots.reserveCapacity(BTREE_MAX_DEPTH)
281+
var childSlots = PackedOffsetList()
284282

285283
var node: _Node = self.root
286284
var startIndex = 0
@@ -292,7 +290,7 @@ extension _BTree {
292290
let endIndex = startIndex + child.read({ $0.numTotalElements })
293291

294292
if offset < endIndex {
295-
childSlots.append(childSlot)
293+
childSlots.append(UInt16(childSlot))
296294
node = child
297295
return nil
298296
} else if offset == endIndex {

‎Sources/SortedCollections/BTree/_BTree.PackedOffsetList.swift‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift Collections open source project
4+
//
5+
// Copyright (c) 2021 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
//
10+
//===----------------------------------------------------------------------===//
11+
12+
@usableFromInline
13+
internal let _PACKED_OFFSET_LIST_MAX_SIZE = 16
14+
15+
extension _BTree {
16+
/// A stack-allocated list of UInt16 values.
17+
@usableFromInline
18+
internal struct PackedOffsetList {
19+
@usableFromInline
20+
internal var depth: UInt8
21+
22+
@usableFromInline
23+
internal var offsets: (
24+
UInt16, UInt16, UInt16, UInt16,
25+
UInt16, UInt16, UInt16, UInt16,
26+
UInt16, UInt16, UInt16, UInt16,
27+
UInt16, UInt16, UInt16, UInt16
28+
)
29+
30+
@inlinable
31+
@inline(__always)
32+
internal init(depth: UInt8 = 0) {
33+
self.depth = depth
34+
self.offsets = (
35+
0, 0, 0, 0,
36+
0, 0, 0, 0,
37+
0, 0, 0, 0,
38+
0, 0, 0, 0
39+
)
40+
}
41+
42+
/// Appends a value to the offset list
43+
@inlinable
44+
@inline(__always)
45+
internal mutating func append(_ offset: UInt16) {
46+
assert(depth < _PACKED_OFFSET_LIST_MAX_SIZE, "Out of bounds access in offset list.")
47+
self.depth &+= 1
48+
self[self.depth] = offset
49+
}
50+
51+
@inlinable
52+
@inline(__always)
53+
internal subscript(_ offset: UInt8) -> UInt16 {
54+
get {
55+
assert(offset <= depth && depth <= _PACKED_OFFSET_LIST_MAX_SIZE, "Out of bounds access in offset list.")
56+
switch offset {
57+
case 0: return self.offsets.0
58+
case 1: return self.offsets.1
59+
case 2: return self.offsets.2
60+
case 3: return self.offsets.3
61+
case 4: return self.offsets.4
62+
case 5: return self.offsets.5
63+
case 6: return self.offsets.6
64+
case 7: return self.offsets.7
65+
case 8: return self.offsets.8
66+
case 9: return self.offsets.9
67+
case 10: return self.offsets.10
68+
case 11: return self.offsets.11
69+
case 12: return self.offsets.12
70+
case 13: return self.offsets.13
71+
case 14: return self.offsets.14
72+
case 15: return self.offsets.15
73+
default: preconditionFailure("Packed offset list too small.")
74+
}
75+
}
76+
77+
_modify {
78+
assert(offset <= depth && depth <= _PACKED_OFFSET_LIST_MAX_SIZE, "Out of bounds access in offset list.")
79+
switch offset {
80+
case 0: yield &self.offsets.0
81+
case 1: yield &self.offsets.1
82+
case 2: yield &self.offsets.2
83+
case 3: yield &self.offsets.3
84+
case 4: yield &self.offsets.4
85+
case 5: yield &self.offsets.5
86+
case 6: yield &self.offsets.6
87+
case 7: yield &self.offsets.7
88+
case 8: yield &self.offsets.8
89+
case 9: yield &self.offsets.9
90+
case 10: yield &self.offsets.10
91+
case 11: yield &self.offsets.11
92+
case 12: yield &self.offsets.12
93+
case 13: yield &self.offsets.13
94+
case 14: yield &self.offsets.14
95+
case 15: yield &self.offsets.15
96+
default: preconditionFailure("Packed offset list too small.")
97+
}
98+
}
99+
}
100+
}
101+
}

‎Sources/SortedCollections/BTree/_BTree.UnsafePath.swift‎

Lines changed: 8 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -20,58 +20,10 @@ extension _BTree {
2020
/// safety, use ``_BTree.Index`` instead
2121
@usableFromInline
2222
internal struct UnsafePath {
23-
@usableFromInline
24-
internal struct PackedOffsetList {
25-
@usableFromInline
26-
internal var depth: UInt8
27-
28-
@usableFromInline
29-
internal var offsets: UInt64
30-
31-
@inlinable
32-
@inline(__always)
33-
internal init() {
34-
self.depth = 0
35-
self.offsets = 0
36-
}
37-
38-
@inlinable
39-
@inline(__always)
40-
internal mutating func pop() {
41-
assert(depth != 0, "Attempted to ascend from root path.")
42-
self.depth &-= 1
43-
}
44-
45-
@inlinable
46-
@inline(__always)
47-
internal mutating func move(by slots: UInt16) {
48-
let level: UInt8 = depth << 4
49-
let mask: UInt64 = UInt64(0xFFFF) << level
50-
51-
var oldSlot = (offsets & mask) >> level
52-
oldSlot &+= UInt64(slots)
53-
54-
offsets = (offsets & ~mask) | (oldSlot << level)
55-
}
56-
57-
@inlinable
58-
@inline(__always)
59-
internal mutating func child(at slot: UInt16) {
60-
assert(depth != 3, "Attempted to exceed maximum depth.")
61-
depth &+= 1
62-
63-
let level: UInt8 = depth << 4
64-
let mask: UInt64 = UInt64(0xFFFF) << level
65-
66-
offsets = (offsets & ~mask) | (UInt64(slot) << level)
67-
}
68-
}
69-
70-
// TODO: potentially make compact (U)Int8/16 type to be more compact
7123
/// The position of each of the parent nodes in their parents. The path's depth
7224
/// is offsets.count + 1
7325
@usableFromInline
74-
internal var childSlots: Array<Int>
26+
internal var childSlots: PackedOffsetList
7527

7628
@usableFromInline
7729
internal unowned(unsafe) var node: Node.Storage
@@ -110,7 +62,7 @@ extension _BTree {
11062
internal init(
11163
node: Node,
11264
slot: Int,
113-
childSlots: Array<Int>,
65+
childSlots: PackedOffsetList,
11466
offset: Int
11567
) {
11668
self.init(node: node.storage, slot: slot, childSlots: childSlots, offset: offset)
@@ -126,7 +78,7 @@ extension _BTree {
12678
internal init(
12779
node: Node.Storage,
12880
slot: Int,
129-
childSlots: Array<Int>,
81+
childSlots: PackedOffsetList,
13082
offset: Int
13183
) {
13284
self.node = node
@@ -164,17 +116,17 @@ extension _BTree.UnsafePath: Comparable {
164116
/// - Complexity: O(`log n`)
165117
@inlinable
166118
public static func <(lhs: _BTree.UnsafePath, rhs: _BTree.UnsafePath) -> Bool {
167-
for i in 0..<min(lhs.childSlots.count, rhs.childSlots.count) {
119+
for i in 0..<min(lhs.childSlots.depth, rhs.childSlots.depth) {
168120
if lhs.childSlots[i] < rhs.childSlots[i] {
169121
return true
170122
}
171123
}
172124

173-
if lhs.childSlots.count < rhs.childSlots.count {
174-
let rhsOffset = rhs.childSlots[lhs.childSlots.count - 1]
125+
if lhs.childSlots.depth < rhs.childSlots.depth {
126+
let rhsOffset = rhs.childSlots[lhs.childSlots.depth - 1]
175127
return lhs.slot < rhsOffset
176-
} else if rhs.childSlots.count < lhs.childSlots.count {
177-
let lhsOffset = lhs.childSlots[rhs.childSlots.count - 1]
128+
} else if rhs.childSlots.depth < lhs.childSlots.depth {
129+
let lhsOffset = lhs.childSlots[rhs.childSlots.depth - 1]
178130
return lhsOffset <= rhs.slot
179131
} else {
180132
return lhs.slot < rhs.slot

‎Sources/SortedCollections/BTree/_BTree.swift‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,11 @@
1414

1515
/// Internal node capacity for BTree
1616
@usableFromInline
17-
internal let BTREE_INTERNAL_CAPACITY = 16
17+
internal let _BTREE_INTERNAL_CAPACITY = 16
1818

1919
/// Leaf node capacity for BTree
2020
@usableFromInline
21-
internal let BTREE_LEAF_CAPACITY = 470
22-
23-
/// An expected rough upper bound for BTree depth
24-
@usableFromInline
25-
internal let BTREE_MAX_DEPTH = 10
21+
internal let _BTREE_LEAF_CAPACITY = 470
2622

2723
/// A bidirectional collection representing a BTree which efficiently stores its
2824
/// elements in sorted order and maintains roughly `O(log count)`
@@ -71,7 +67,7 @@ internal struct _BTree<Key: Comparable, Value> {
7167
/// - internalCapacity: The capacity of the internal nodes. Generally prefered to be less than `leafCapacity`.
7268
@inlinable
7369
@inline(__always)
74-
internal init(leafCapacity: Int = BTREE_LEAF_CAPACITY, internalCapacity: Int = BTREE_INTERNAL_CAPACITY) {
70+
internal init(leafCapacity: Int = _BTREE_LEAF_CAPACITY, internalCapacity: Int = _BTREE_INTERNAL_CAPACITY) {
7571
self.init(
7672
rootedAt: Node(withCapacity: leafCapacity, isLeaf: true),
7773
leafCapacity: leafCapacity,
@@ -100,7 +96,7 @@ internal struct _BTree<Key: Comparable, Value> {
10096
/// - internalCapacity: The capacity of the internal nodes. Generally prefered to be less than `leafCapacity`.
10197
@inlinable
10298
@inline(__always)
103-
internal init(rootedAt root: Node, leafCapacity: Int = BTREE_LEAF_CAPACITY, internalCapacity: Int = BTREE_INTERNAL_CAPACITY) {
99+
internal init(rootedAt root: Node, leafCapacity: Int = _BTREE_LEAF_CAPACITY, internalCapacity: Int = _BTREE_INTERNAL_CAPACITY) {
104100
self.root = root
105101
self.internalCapacity = internalCapacity
106102
self.age = Int32(truncatingIfNeeded: ObjectIdentifier(root.storage).hashValue)

0 commit comments

Comments
 (0)