Fix deletion when node4/16 is full

apple · Aug 19, 2023 · f9a5452 · f9a5452
1 parent 91e2c7c
commit f9a5452
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 11 deletions.
diff --git a/Sources/ARTreeModule/Node16.swift b/Sources/ARTreeModule/Node16.swift
@@ -194,12 +194,9 @@ extension Node16: InternalNode {
     assert(index < Self.numKeys, "index can't >= 16 in Node16")
     assert(index < count, "not enough childs in node")
 
-    let childBuf = child(at: index)
-    // childBuf?.deallocate()
-
     withBody { keys, childs in
-      keys[self.count] = 0
-      childs[self.count] = nil
+      keys[index] = 0
+      childs[index] = nil
 
       count -= 1
       keys.shiftLeft(startIndex: index + 1, endIndex: count, by: 1)

diff --git a/Sources/ARTreeModule/Node256.swift b/Sources/ARTreeModule/Node256.swift
@@ -130,7 +130,6 @@ extension Node256: InternalNode {
 
   public mutating func deleteChild(at index: Index, ref: ChildSlotPtr?) {
     return withBody { childs in
-      // childs[index]?.deallocate()
       childs[index] = nil
       count -= 1
 

diff --git a/Sources/ARTreeModule/Node4.swift b/Sources/ARTreeModule/Node4.swift
@@ -168,12 +168,9 @@ extension Node4: InternalNode {
     assert(index < 4, "index can't >= 4 in Node4")
     assert(index < count, "not enough childs in node")
 
-    let childBuf = child(at: index)
-    // childBuf?.deallocate()
-
     withBody { keys, childs in
-      keys[self.count] = 0
-      childs[self.count] = nil
+      keys[index] = 0
+      childs[index] = nil
 
       count -= 1
       keys.shiftLeft(startIndex: index + 1, endIndex: count, by: 1)

diff --git a/Tests/ARTreeModuleTests/Node4Tests.swift b/Tests/ARTreeModuleTests/Node4Tests.swift
@@ -75,6 +75,35 @@ final class ARTreeNode4Tests: XCTestCase {
     XCTAssertEqual(ptr?.type, .leaf)
   }
 
+  func test4DeleteFromFull() throws {
+    var node = Node4.allocate()
+    node.addChild(forKey: 1, node: NodeLeaf.allocate(key: [1], value: 1, of: Int.self))
+    node.addChild(forKey: 2, node: NodeLeaf.allocate(key: [2], value: 2, of: Int.self))
+    node.addChild(forKey: 3, node: NodeLeaf.allocate(key: [3], value: 3, of: Int.self))
+    node.addChild(forKey: 4, node: NodeLeaf.allocate(key: [4], value: 4, of: Int.self))
+    XCTAssertEqual(node.type, .node4)
+    XCTAssertEqual(
+      node.print(value: Int.self),
+      "○ Node4 {childs=4, partial=[]}\n" +
+      "├──○ 1: 1[1] -> 1\n" +
+      "├──○ 2: 1[2] -> 2\n" +
+      "├──○ 3: 1[3] -> 3\n" +
+      "└──○ 4: 1[4] -> 4")
+    node.deleteChild(at: 1)
+    XCTAssertEqual(
+      node.print(value: Int.self),
+      "○ Node4 {childs=3, partial=[]}\n" +
+      "├──○ 1: 1[1] -> 1\n" +
+      "├──○ 3: 1[3] -> 3\n" +
+      "└──○ 4: 1[4] -> 4")
+
+    var ptr: RawNode? = node.rawNode
+    node.deleteChild(at: 1, ref: &ptr)
+    node.deleteChild(at: 1, ref: &ptr)
+    XCTAssertNotNil(ptr)
+    XCTAssertEqual(ptr?.type, .leaf)
+  }
+
   func test4ExapandTo16() throws {
     var node = Node4.allocate()
     node.addChild(forKey: 1, node: NodeLeaf.allocate(key: [1], value: [1], of: [UInt8].self))