From 5fa86d54ddbdb16e9c824c0cee83a314fe450d69 Mon Sep 17 00:00:00 2001 From: Vishesh Yadav Date: Thu, 24 Aug 2023 18:56:24 -0700 Subject: [PATCH] Refactor to use getters --- Sources/ARTreeModule/ARTree+insert.swift | 2 +- Sources/ARTreeModule/InternalNode.swift | 2 +- .../ARTreeModule/Node+StringConverter.swift | 93 ++++--- Sources/ARTreeModule/Node16.swift | 193 ++++++-------- Sources/ARTreeModule/Node256.swift | 126 ++++------ Sources/ARTreeModule/Node4.swift | 164 ++++++------ Sources/ARTreeModule/Node48.swift | 237 ++++++++---------- Tests/ARTreeModuleTests/DeleteTests.swift | 1 + 8 files changed, 349 insertions(+), 469 deletions(-) diff --git a/Sources/ARTreeModule/ARTree+insert.swift b/Sources/ARTreeModule/ARTree+insert.swift index 2f3c7509a..26c103312 100644 --- a/Sources/ARTreeModule/ARTree+insert.swift +++ b/Sources/ARTreeModule/ARTree+insert.swift @@ -54,7 +54,7 @@ extension ARTree { var next = Node4.allocate() _ = next.addChild(forKey: key[start - 1], node: newNode) newNode = next - longestPrefix -= 1 // One keys goes for mapping the child in next node. + longestPrefix -= 1 // One keys goes for mapping the child in next node. } ref.pointee = newNode.rawNode // Replace child in parent. diff --git a/Sources/ARTreeModule/InternalNode.swift b/Sources/ARTreeModule/InternalNode.swift index aedea6019..d37f9a26c 100644 --- a/Sources/ARTreeModule/InternalNode.swift +++ b/Sources/ARTreeModule/InternalNode.swift @@ -218,7 +218,7 @@ extension InternalNode { let slf: Node4 = selfRef as! Node4 var node: any InternalNode = newValue.toInternalNode() node.partialBytes.shiftRight() - node.partialBytes[0] = slf.withBody { k, _ in k[0] } + node.partialBytes[0] = slf.keys[0] node.partialLength += 1 } return .replaceWith(newValue) diff --git a/Sources/ARTreeModule/Node+StringConverter.swift b/Sources/ARTreeModule/Node+StringConverter.swift index 40d722127..c11cc4208 100644 --- a/Sources/ARTreeModule/Node+StringConverter.swift +++ b/Sources/ARTreeModule/Node+StringConverter.swift @@ -111,16 +111,15 @@ extension Node4: NodePrettyPrinter { func prettyPrint(depth: Int) -> String { let addr = Const.testPrintAddr ? " \(_addressString(for: self.rawNode.buf))" : " " var output = "Node4\(addr){childs=\(count), partial=\(partial)}\n" - withBody { keys, childs in - for idx in 0.. String { let addr = Const.testPrintAddr ? " \(_addressString(for: self.rawNode.buf))" : " " var output = "Node16\(addr){childs=\(count), partial=\(partial)}\n" - withBody { keys, childs in - for idx in 0..= 0xFF { - continue - } - - total += 1 - let last = total == count - output += indent(depth, last: last) - output += String(key) + ": " - output += child(at: Int(slot))!.prettyPrint(depth: depth + 1, with: Spec.self) - if !last { - output += "\n" - } + + for (key, slot) in keys.enumerated() { + if slot >= 0xFF { + continue + } + + total += 1 + let last = total == count + output += indent(depth, last: last) + output += String(key) + ": " + output += child(at: Int(slot))!.prettyPrint(depth: depth + 1, with: Spec.self) + if !last { + output += "\n" } } @@ -183,22 +180,22 @@ extension Node256: NodePrettyPrinter { let addr = Const.testPrintAddr ? " \(_addressString(for: self.rawNode.buf))" : " " var output = "Node256\(addr){childs=\(count), partial=\(partial)}\n" var total = 0 - withBody { childs in - for (key, child) in childs.enumerated() { - if child == nil { - continue - } - - total += 1 - let last = total == count - output += indent(depth, last: last) - output += String(key) + ": " - output += child!.prettyPrint(depth: depth + 1, with: Spec.self) - if !last { - output += "\n" - } + + for (key, child) in childs.enumerated() { + if child == nil { + continue + } + + total += 1 + let last = total == count + output += indent(depth, last: last) + output += String(key) + ": " + output += child!.prettyPrint(depth: depth + 1, with: Spec.self) + if !last { + output += "\n" } } + return output } } diff --git a/Sources/ARTreeModule/Node16.swift b/Sources/ARTreeModule/Node16.swift index 2a13c53c8..75347c2ee 100644 --- a/Sources/ARTreeModule/Node16.swift +++ b/Sources/ARTreeModule/Node16.swift @@ -18,17 +18,34 @@ extension Node16 { static var numKeys: Int { 16 } } +extension Node16 { + var keys: UnsafeMutableBufferPointer { + storage.withBodyPointer { + UnsafeMutableBufferPointer( + start: $0.assumingMemoryBound(to: KeyPart.self), + count: Self.numKeys + ) + } + } + + var childs: UnsafeMutableBufferPointer { + storage.withBodyPointer { + let childPtr = $0.advanced(by: Self.numKeys * MemoryLayout.stride) + .assumingMemoryBound(to: RawNode?.self) + return UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + } + } +} + extension Node16 { static func allocate() -> NodeStorage { let storage = NodeStorage.allocate() storage.update { node in - node.withBody { keys, childs in - UnsafeMutableRawPointer(keys.baseAddress!) - .bindMemory(to: UInt8.self, capacity: Self.numKeys) - UnsafeMutableRawPointer(childs.baseAddress!) - .bindMemory(to: RawNode?.self, capacity: Self.numKeys) - } + UnsafeMutableRawPointer(node.keys.baseAddress!) + .bindMemory(to: UInt8.self, capacity: Self.numKeys) + UnsafeMutableRawPointer(node.childs.baseAddress!) + .bindMemory(to: RawNode?.self, capacity: Self.numKeys) } return storage @@ -37,16 +54,12 @@ extension Node16 { static func allocate(copyFrom: Node4) -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: copyFrom) - copyFrom.withBody { fromKeys, fromChilds in - node.withBody { newKeys, newChilds in - UnsafeMutableRawBufferPointer(newKeys).copyBytes(from: fromKeys) - UnsafeMutableRawBufferPointer(newChilds).copyBytes( - from: UnsafeMutableRawBufferPointer(fromChilds)) - Self.retainChildren(newChilds, count: node.count) - } - } + storage.update { newNode in + newNode.copyHeader(from: copyFrom) + UnsafeMutableRawBufferPointer(newNode.keys).copyBytes(from: copyFrom.keys) + UnsafeMutableRawBufferPointer(newNode.childs).copyBytes( + from: UnsafeMutableRawBufferPointer(copyFrom.childs)) + Self.retainChildren(newNode.childs, count: newNode.count) } return storage @@ -55,53 +68,29 @@ extension Node16 { static func allocate(copyFrom: Node48) -> NodeStorage { let storage = NodeStorage.allocate() - storage.update { node in - node.copyHeader(from: copyFrom) - copyFrom.withBody { fromKeys, fromChilds in - node.withBody { newKeys, newChilds in - var slot = 0 - for key: UInt8 in 0...255 { - let childPosition = Int(fromKeys[Int(key)]) - if childPosition == 0xFF { - continue - } - - newKeys[slot] = key - newChilds[slot] = fromChilds[childPosition] - slot += 1 - } - - assert(slot == node.count) - Self.retainChildren(newChilds, count: node.count) + storage.update { newNode in + newNode.copyHeader(from: copyFrom) + + var slot = 0 + for key: UInt8 in 0...255 { + let childPosition = Int(copyFrom.keys[Int(key)]) + if childPosition == 0xFF { + continue } + + newNode.keys[slot] = key + newNode.childs[slot] = copyFrom.childs[childPosition] + slot += 1 } + + assert(slot == newNode.count) + Self.retainChildren(newNode.childs, count: newNode.count) } return storage } } -extension Node16 { - typealias Keys = UnsafeMutableBufferPointer - typealias Children = UnsafeMutableBufferPointer - - func withBody(body: (Keys, Children) 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 @@ -109,15 +98,13 @@ extension Node16: InternalNode { } func index(forKey k: KeyPart) -> Index? { - return withBody { keys, _ in - for (index, key) in keys.enumerated() { - if key == k { - return index - } + for (index, key) in keys.enumerated() { + if key == k { + return index } - - return nil } + + return nil } func index() -> Index? { @@ -135,34 +122,28 @@ extension Node16: InternalNode { return nil } - return withBody { keys, _ in - for idx in 0..= Int(k) { - return idx - } + for idx in 0..= Int(k) { + return idx } - - return count } + + return count } func child(at: Index) -> RawNode? { assert(at < Self.numKeys, "maximum \(Self.numKeys) childs allowed") - return withBody { _, childs in - return childs[at] - } + return childs[at] } mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult { if let slot = _insertSlot(forKey: k) { - withBody { keys, childs in - assert(count == 0 || keys[slot] != k, "node for key \(k) already exists") - keys.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) - childs.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) - keys[slot] = k - childs[slot] = node - count += 1 - } + assert(count == 0 || keys[slot] != k, "node for key \(k) already exists") + keys.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) + childs.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) + keys[slot] = k + childs[slot] = node + count += 1 return .noop } else { return Node48.allocate(copyFrom: self).update { newNode in @@ -176,31 +157,27 @@ extension Node16: InternalNode { assert(index < Self.numKeys, "index can't >= 16 in Node16") assert(index < count, "not enough childs in node") - return withBody { keys, childs in - keys[index] = 0 - childs[index] = nil + keys[index] = 0 + childs[index] = nil - count -= 1 - keys.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) - childs.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) - childs[count] = nil // Clear the last item. + count -= 1 + keys.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) + childs.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) + childs[count] = nil // Clear the last item. - if count == 3 { - // Shrink to Node4. - let newNode = Node4.allocate(copyFrom: self) - return .replaceWith(newNode.node.rawNode) - } - - return .noop + if count == 3 { + // Shrink to Node4. + let newNode = Node4.allocate(copyFrom: self) + return .replaceWith(newNode.node.rawNode) } + + return .noop } mutating func withChildRef(at index: Index, _ body: (RawNode.SlotRef) -> R) -> R { assert(index < count, "not enough childs in node") - return withBody { _, childs in - let ref = childs.baseAddress! + index - return body(ref) - } + let ref = childs.baseAddress! + index + return body(ref) } } @@ -208,10 +185,8 @@ extension Node16: ArtNode { final class Buffer: RawNodeBuffer { deinit { var node = Node16(buffer: self) - node.withBody { _, childs in - for idx in 0..<16 { - childs[idx] = nil - } + for idx in 0..<16 { + node.childs[idx] = nil } node.count = 0 } @@ -220,15 +195,11 @@ extension Node16: ArtNode { func clone() -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: self) - self.withBody { fromKeys, fromChildren in - node.withBody { newKeys, newChildren in - for idx in 0.. { + storage.withBodyPointer { + UnsafeMutableBufferPointer( + start: $0.assumingMemoryBound(to: RawNode?.self), + count: 256) + } + } +} + extension Node256 { static func allocate() -> NodeStorage { let storage = NodeStorage.allocate() - storage.update { node in - _ = node.withBody { childs in - UnsafeMutableRawPointer(childs.baseAddress!) - .bindMemory(to: RawNode?.self, capacity: Self.numKeys) - } + storage.update { newNode in + UnsafeMutableRawPointer(newNode.childs.baseAddress!) + .bindMemory(to: RawNode?.self, capacity: Self.numKeys) } return storage @@ -35,50 +43,30 @@ extension Node256 { static func allocate(copyFrom: Node48) -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: copyFrom) - copyFrom.withBody { fromKeys, fromChilds in - node.withBody { newChilds in - for key in 0..<256 { - let slot = Int(fromKeys[key]) - if slot < 0xFF { - newChilds[key] = fromChilds[slot] - } - } - - Self.retainChildren(newChilds, count: Self.numKeys) + storage.update { newNode in + newNode.copyHeader(from: copyFrom) + for key in 0..<256 { + let slot = Int(copyFrom.keys[key]) + if slot < 0xFF { + newNode.childs[key] = copyFrom.childs[slot] } } - assert(node.count == 48, "should have exactly 48 childs") + + Self.retainChildren(newNode.childs, count: Self.numKeys) + assert(newNode.count == 48, "should have exactly 48 childs") } return storage } } -extension Node256 { - typealias Keys = UnsafeMutableBufferPointer - typealias Children = UnsafeMutableBufferPointer - - func withBody(body: (Children) 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 } func index(forKey k: KeyPart) -> Index? { - return withBody { childs in - return childs[Int(k)] != nil ? Int(k) : nil - } + return childs[Int(k)] != nil ? Int(k) : nil } func index() -> Index? { @@ -86,53 +74,43 @@ extension Node256: InternalNode { } func next(index: Index) -> Index? { - return withBody { childs in - for idx in index + 1..<256 { - if childs[idx] != nil { - return idx - } + for idx in index + 1..<256 { + if childs[idx] != nil { + return idx } - - return nil } + + return nil } func child(at: Int) -> RawNode? { assert(at < 256, "maximum 256 childs allowed") - return withBody { childs in - return childs[at] - } + return childs[at] } mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult { - return withBody { childs in - assert(childs[Int(k)] == nil, "node for key \(k) already exists") - childs[Int(k)] = node - count += 1 - return .noop - } + assert(childs[Int(k)] == nil, "node for key \(k) already exists") + childs[Int(k)] = node + count += 1 + return .noop } - public mutating func deleteChild(at index: Index) -> UpdateResult { - return withBody { childs in - childs[index] = nil - count -= 1 + mutating func deleteChild(at index: Index) -> UpdateResult { + childs[index] = nil + count -= 1 - if count == 40 { - let newNode = Node48.allocate(copyFrom: self) - return .replaceWith(newNode.node.rawNode) - } - - return .noop + if count == 40 { + let newNode = Node48.allocate(copyFrom: self) + return .replaceWith(newNode.node.rawNode) } + + return .noop } mutating func withChildRef(at index: Index, _ body: (RawNode.SlotRef) -> R) -> R { assert(index < count, "not enough childs in node") - return withBody { childs in - let ref = childs.baseAddress! + index - return body(ref) - } + let ref = childs.baseAddress! + index + return body(ref) } } @@ -140,10 +118,8 @@ extension Node256: ArtNode { final class Buffer: RawNodeBuffer { deinit { var node = Node256(buffer: self) - node.withBody { childs in - for idx in 0..<256 { - childs[idx] = nil - } + for idx in 0..<256 { + node.childs[idx] = nil } node.count = 0 } @@ -152,14 +128,10 @@ extension Node256: ArtNode { func clone() -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: self) - self.withBody { fromChildren in - node.withBody { newChildren in - for idx in 0..<256 { - newChildren[idx] = fromChildren[idx] - } - } + storage.update { newNode in + newNode.copyHeader(from: self) + for idx in 0..<256 { + newNode.childs[idx] = childs[idx] } } diff --git a/Sources/ARTreeModule/Node4.swift b/Sources/ARTreeModule/Node4.swift index ce603b1c2..ac7434846 100644 --- a/Sources/ARTreeModule/Node4.swift +++ b/Sources/ARTreeModule/Node4.swift @@ -18,17 +18,34 @@ extension Node4 { static var numKeys: Int { 4 } } +extension Node4 { + var keys: UnsafeMutableBufferPointer { + storage.withBodyPointer { + UnsafeMutableBufferPointer( + start: $0.assumingMemoryBound(to: KeyPart.self), + count: Self.numKeys + ) + } + } + + var childs: UnsafeMutableBufferPointer { + storage.withBodyPointer { + let childPtr = $0.advanced(by: Self.numKeys * MemoryLayout.stride) + .assumingMemoryBound(to: RawNode?.self) + return UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + } + } +} + extension Node4 { static func allocate() -> NodeStorage { let storage = NodeStorage.allocate() storage.update { node in - node.withBody { keys, childs in - UnsafeMutableRawPointer(keys.baseAddress!) - .bindMemory(to: UInt8.self, capacity: Self.numKeys) - UnsafeMutableRawPointer(childs.baseAddress!) - .bindMemory(to: RawNode?.self, capacity: Self.numKeys) - } + UnsafeMutableRawPointer(node.keys.baseAddress!) + .bindMemory(to: UInt8.self, capacity: Self.numKeys) + UnsafeMutableRawPointer(node.childs.baseAddress!) + .bindMemory(to: RawNode?.self, capacity: Self.numKeys) } return storage @@ -37,42 +54,19 @@ extension Node4 { static func allocate(copyFrom: Node16) -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: copyFrom) - node.withBody { newKeys, newChilds in - copyFrom.withBody { fromKeys, fromChilds in - UnsafeMutableRawBufferPointer(newKeys).copyBytes( - from: UnsafeBufferPointer(rebasing: fromKeys[0.. - typealias Children = UnsafeMutableBufferPointer + UnsafeMutableRawBufferPointer(newNode.keys).copyBytes( + from: UnsafeBufferPointer(rebasing: copyFrom.keys[0..(body: (Keys, Children) 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) + Self.retainChildren(newNode.childs, count: newNode.count) } + + return storage } } @@ -83,15 +77,13 @@ extension Node4: InternalNode { } func index(forKey k: KeyPart) -> Index? { - return withBody { keys, _ in - for (index, key) in keys.enumerated() { - if key == k { - return index - } + for (index, key) in keys.enumerated() { + if key == k { + return index } - - return nil } + + return nil } func index() -> Index? { @@ -108,34 +100,28 @@ extension Node4: InternalNode { return nil } - return withBody { keys, _ in - for idx in 0..= Int(k) { - return idx - } + for idx in 0..= Int(k) { + return idx } - - return count } + + return count } func child(at: Index) -> RawNode? { assert(at < Self.numKeys, "maximum \(Self.numKeys) childs allowed, given index = \(at)") - return withBody { _, childs in - return childs[at] - } + return childs[at] } mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult { if let slot = _insertSlot(forKey: k) { - withBody { keys, childs in - assert(count == 0 || keys[slot] != k, "node for key \(k) already exists") - keys.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) - childs.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) - keys[slot] = k - childs[slot] = node - count += 1 - } + assert(count == 0 || keys[slot] != k, "node for key \(k) already exists") + keys.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) + childs.shiftRight(startIndex: slot, endIndex: count - 1, by: 1) + keys[slot] = k + childs[slot] = node + count += 1 return .noop } else { var newNode = Node16.allocate(copyFrom: self) @@ -148,30 +134,26 @@ extension Node4: InternalNode { assert(index < 4, "index can't >= 4 in Node4") assert(index < count, "not enough childs in node") - return withBody { keys, childs in - keys[index] = 0 - childs[index] = nil + keys[index] = 0 + childs[index] = nil - count -= 1 - keys.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) - childs.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) - childs[count] = nil // Clear the last item. - - if count == 1 { - // Shrink to leaf node. - return .replaceWith(childs[0]) - } + count -= 1 + keys.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) + childs.shiftLeft(startIndex: index + 1, endIndex: count, by: 1) + childs[count] = nil // Clear the last item. - return .noop + if count == 1 { + // Shrink to leaf node. + return .replaceWith(childs[0]) } + + return .noop } mutating func withChildRef(at index: Index, _ body: (RawNode.SlotRef) -> R) -> R { assert(index < count, "index=\(index) less than count=\(count)") - return withBody { _, childs in - let ref = childs.baseAddress! + index - return body(ref) - } + let ref = childs.baseAddress! + index + return body(ref) } } @@ -179,10 +161,8 @@ extension Node4: ArtNode { final class Buffer: RawNodeBuffer { deinit { var node = Node4(buffer: self) - node.withBody { _, childs in - for idx in 0..<4 { - childs[idx] = nil - } + for idx in 0..<4 { + node.childs[idx] = nil } node.count = 0 } @@ -191,15 +171,11 @@ extension Node4: ArtNode { func clone() -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: self) - self.withBody { fromKeys, fromChildren in - node.withBody { newKeys, newChildren in - for idx in 0.. { + storage.withBodyPointer { + UnsafeMutableBufferPointer( + start: $0.assumingMemoryBound(to: KeyPart.self), + count: 256 + ) + } + } + + var childs: UnsafeMutableBufferPointer { + storage.withBodyPointer { + let childPtr = $0.advanced(by: 256 * MemoryLayout.stride) + .assumingMemoryBound(to: RawNode?.self) + return UnsafeMutableBufferPointer(start: childPtr, count: Self.numKeys) + } + } +} + extension Node48 { static func allocate() -> NodeStorage { let storage = NodeStorage.allocate() - storage.update { node in - node.withBody { keys, childs in - UnsafeMutableRawPointer(keys.baseAddress!) - .bindMemory(to: UInt8.self, capacity: Self.numKeys) - UnsafeMutableRawPointer(childs.baseAddress!) - .bindMemory(to: RawNode?.self, capacity: Self.numKeys) + storage.update { newNode in + UnsafeMutableRawPointer(newNode.keys.baseAddress!) + .bindMemory(to: UInt8.self, capacity: Self.numKeys) + UnsafeMutableRawPointer(newNode.childs.baseAddress!) + .bindMemory(to: RawNode?.self, capacity: Self.numKeys) - for idx in 0..<256 { - keys[idx] = 0xFF - } + for idx in 0..<256 { + newNode.keys[idx] = 0xFF } } @@ -41,19 +58,15 @@ extension Node48 { static func allocate(copyFrom: Node16) -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: copyFrom) - copyFrom.withBody { fromKeys, fromChilds in - node.withBody { newKeys, newChilds in - UnsafeMutableRawBufferPointer(newChilds).copyBytes( - from: UnsafeMutableRawBufferPointer(fromChilds)) - for (idx, key) in fromKeys.enumerated() { - newKeys[Int(key)] = UInt8(idx) - } - - Self.retainChildren(newChilds, count: node.count) - } + storage.update { newNode in + newNode.copyHeader(from: copyFrom) + UnsafeMutableRawBufferPointer(newNode.childs).copyBytes( + from: UnsafeMutableRawBufferPointer(copyFrom.childs)) + for (idx, key) in copyFrom.keys.enumerated() { + newNode.keys[Int(key)] = UInt8(idx) } + + Self.retainChildren(newNode.childs, count: newNode.count) } return storage @@ -62,55 +75,26 @@ extension Node48 { static func allocate(copyFrom: Node256) -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: copyFrom) - copyFrom.withBody { fromChilds in - node.withBody { newKeys, newChilds in - var slot = 0 - for (key, child) in fromChilds.enumerated() { - if child == nil { - continue - } - - newKeys[key] = UInt8(slot) - newChilds[slot] = child - slot += 1 - } - - Self.retainChildren(newChilds, count: node.count) + storage.update { newNode in + newNode.copyHeader(from: copyFrom) + var slot = 0 + for (key, child) in copyFrom.childs.enumerated() { + if child == nil { + continue } + + newNode.keys[key] = UInt8(slot) + newNode.childs[slot] = child + slot += 1 } + + Self.retainChildren(newNode.childs, count: newNode.count) } return storage } } -extension Node48 { - typealias Keys = UnsafeMutableBufferPointer - typealias Children = UnsafeMutableBufferPointer - - func withBody(body: (Keys, Children) 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) - - return try body(keys, childs) - } - } -} - extension Node48: InternalNode { static var size: Int { MemoryLayout.stride + 256 * MemoryLayout.stride + Self.numKeys @@ -118,10 +102,8 @@ extension Node48: InternalNode { } func index(forKey k: KeyPart) -> Index? { - return withBody { keys, _ in - let childIndex = Int(keys[Int(k)]) - return childIndex == 0xFF ? nil : childIndex - } + let childIndex = Int(keys[Int(k)]) + return childIndex == 0xFF ? nil : childIndex } func index() -> Index? { @@ -129,38 +111,31 @@ extension Node48: InternalNode { } func next(index: Index) -> Index? { - return withBody { keys, _ in - for idx: Int in index + 1..<256 { - if keys[idx] != 0xFF { - return Int(keys[idx]) - } + for idx: Int in index + 1..<256 { + if keys[idx] != 0xFF { + return Int(keys[idx]) } - - return nil } + + return nil } func child(at: Int) -> RawNode? { assert(at < Self.numKeys, "maximum \(Self.numKeys) childs allowed") - return withBody { _, childs in - return childs[at] - } + return childs[at] } mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult { if count < Self.numKeys { - withBody { keys, childs in - assert(keys[Int(k)] == 0xFF, "node for key \(k) already exists") + assert(keys[Int(k)] == 0xFF, "node for key \(k) already exists") - guard let slot = findFreeSlot() else { - assert(false, "cannot find free slot in Node48") - return - } - - keys[Int(k)] = KeyPart(slot) - childs[slot] = node + guard let slot = findFreeSlot() else { + fatalError("cannot find free slot in Node48") } + keys[Int(k)] = KeyPart(slot) + childs[slot] = node + self.count += 1 return .noop } else { @@ -172,59 +147,53 @@ extension Node48: InternalNode { } public mutating func deleteChild(at index: Index) -> UpdateResult { - return withBody { keys, childs in - let targetSlot = Int(keys[index]) - assert(targetSlot != 0xFF, "slot is empty already") - // 1. Find out who has the last slot. - var lastSlotKey = 0 - for k in 0..<256 { - if keys[k] == count - 1 { - lastSlotKey = k - break - } + let targetSlot = Int(keys[index]) + assert(targetSlot != 0xFF, "slot is empty already") + // 1. Find out who has the last slot. + var lastSlotKey = 0 + for k in 0..<256 { + if keys[k] == count - 1 { + lastSlotKey = k + break } + } - // 2. Move last child slot into current child slot, and reset last child slot. - childs[targetSlot] = childs[count - 1] - childs[count - 1] = nil - - // 3. Map that key to current slot. - keys[lastSlotKey] = UInt8(targetSlot) + // 2. Move last child slot into current child slot, and reset last child slot. + childs[targetSlot] = childs[count - 1] + childs[count - 1] = nil - // 4. Clear input key. - keys[index] = 0xFF + // 3. Map that key to current slot. + keys[lastSlotKey] = UInt8(targetSlot) - // 5. Reduce number of children. - count -= 1 + // 4. Clear input key. + keys[index] = 0xFF - // 6. Shrink the node to Node16 if needed. - if count == 13 { - let newNode = Node16.allocate(copyFrom: self) - return .replaceWith(newNode.node.rawNode) - } + // 5. Reduce number of children. + count -= 1 - return .noop + // 6. Shrink the node to Node16 if needed. + if count == 13 { + let newNode = Node16.allocate(copyFrom: self) + return .replaceWith(newNode.node.rawNode) } + + return .noop } private func findFreeSlot() -> Int? { - return withBody { _, childs in - for (index, child) in childs.enumerated() { - if child == nil { - return index - } + for (index, child) in childs.enumerated() { + if child == nil { + return index } - - return nil } + + return nil } mutating func withChildRef(at index: Index, _ body: (RawNode.SlotRef) -> R) -> R { assert(index < count, "not enough childs in node") - return withBody { _, childs in - let ref = childs.baseAddress! + index - return body(ref) - } + let ref = childs.baseAddress! + index + return body(ref) } } @@ -232,10 +201,8 @@ extension Node48: ArtNode { final class Buffer: RawNodeBuffer { deinit { var node = Node48(buffer: self) - node.withBody { _, childs in - for idx in 0..<48 { - childs[idx] = nil - } + for idx in 0..<48 { + node.childs[idx] = nil } node.count = 0 } @@ -244,17 +211,13 @@ extension Node48: ArtNode { func clone() -> NodeStorage { let storage = Self.allocate() - storage.update { node in - node.copyHeader(from: self) - self.withBody { fromKeys, fromChildren in - node.withBody { newKeys, newChildren in - for idx in 0..<256 { - let slot = fromKeys[idx] - newKeys[idx] = slot - if slot != 0xFF { - newChildren[Int(slot)] = fromChildren[Int(slot)] - } - } + storage.update { newNode in + newNode.copyHeader(from: self) + for idx in 0..<256 { + let slot = keys[idx] + newNode.keys[idx] = slot + if slot != 0xFF { + newNode.childs[Int(slot)] = childs[Int(slot)] } } } diff --git a/Tests/ARTreeModuleTests/DeleteTests.swift b/Tests/ARTreeModuleTests/DeleteTests.swift index 6b56cec86..6fd47df9e 100644 --- a/Tests/ARTreeModuleTests/DeleteTests.swift +++ b/Tests/ARTreeModuleTests/DeleteTests.swift @@ -216,6 +216,7 @@ final class ARTreeDeleteTests: XCTestCase { t.delete(key: [i, i + 1]) } } + XCTAssertEqual(t._root?.type, .node48) XCTAssertEqual( t.description, "○ Node48 {childs=40, partial=[]}\n" +