From 0ee791db3c9912229a7400a4518ea003e2b2caf5 Mon Sep 17 00:00:00 2001 From: Janis Date: Mon, 20 Mar 2023 01:15:18 +0100 Subject: [PATCH] removal works better but still somehow broken somewhere --- src/main.zig | 156 ++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 131 insertions(+), 25 deletions(-) diff --git a/src/main.zig b/src/main.zig index 3fb454c..25a1a6f 100644 --- a/src/main.zig +++ b/src/main.zig @@ -55,7 +55,7 @@ const Optional = struct { return !self.is_some(); } - pub fn is_unwrap(self: Self) T { + pub fn unwrap(self: Self) T { return self.Some; } @@ -191,6 +191,19 @@ const BTree = struct { } } + fn remove(self: *Self, key: u32) void { + if (self.root) |*root| { + const search = root.find_key(key); + + switch (search) { + .Leaf => |leaf| { + NodeOrLeaf.from_leaf(leaf.leaf).remove_key_from_node(key); + }, + .Edge => {}, + } + } + } + fn find_key(self: *Self, key: u32) ?u32 { switch (self.root.?.find_key(key)) { .Leaf => |leaf| { @@ -279,7 +292,7 @@ const BTree = struct { for (node.edges, 0..) |edge, i| { if (edge) |edge_node| { - edge_node.as_leaf().parent = .{ .parent = node, .idx = @intCast(u16, i) }; + edge_node.as_leaf().set_parent(node, @intCast(u16, i)); } } @@ -362,23 +375,26 @@ const BTree = struct { const right = Optional.into_option(node.get_edge(index + 1)); if (right.is_some() and left.is_some()) { - const right_value = right.unwrap().find_least_significant_key(); - const left_value = left.unwrap().find_most_significant_key(); + const right_value = right.Some.find_least_significant_key(); + const left_value = left.Some.find_most_significant_key(); // pick the leaf with more values to reduce unneccesairy rebalances - const myleaf = if (right_value.leaf.len > left_value.leaf.len) { - // use right leaf - right_value; - } else { - // use left leaf - left_value; + const myleaf = blk: { + if (right_value.leaf.len > left_value.leaf.len) { + // use right leaf + break :blk right_value; + } else { + // use left leaf + break :blk left_value; + } }; - const promoted = myleaf.get_values()[myleaf.idx]; - myleaf.remove_key(promoted); - self.leaf.values[index] = promoted; + const promoted = myleaf.leaf.get_values()[myleaf.idx]; + myleaf.leaf.remove_key(promoted); + self.as_leaf().values[index] = promoted; // TODO: check `myleaf` for rebalance + myleaf.leaf.rebalance(); } else { // no adjacent edges // (I don't think this can actually happen in a proper rtree) @@ -389,6 +405,7 @@ const BTree = struct { // instead we need to check for rebalance of self here TODO node.shift_edges_left(index, 1); + self.as_leaf().rebalance(); } }, .leaf => |node| { @@ -397,9 +414,7 @@ const BTree = struct { node.remove_key(key); - if (self.as_leaf().needs_rebalance()) { - std.debug.print("---------- require rebalance! -----------\n", .{}); - } + self.as_leaf().rebalance(); }, } } @@ -435,7 +450,12 @@ const BTree = struct { } } - fn find_most_significant_key(self: NodeOrLeaf) struct { leaf: *Leaf, idx: u16 } { + const SignificantKeyResult = struct { + leaf: *Leaf, + idx: u16, + }; + + fn find_most_significant_key(self: NodeOrLeaf) SignificantKeyResult { var node = self; while (true) { switch (node.force()) { @@ -452,7 +472,7 @@ const BTree = struct { } } - fn find_least_significant_key(self: NodeOrLeaf) struct { leaf: *Leaf, idx: u16 } { + fn find_least_significant_key(self: NodeOrLeaf) SignificantKeyResult { var node = self; while (true) { switch (node.force()) { @@ -559,7 +579,7 @@ const BTree = struct { edge.dbg(); std.debug.print("\n", .{}); } else { - child.as_leaf().parent = .{ .parent = self, .idx = idx }; + child.as_leaf().set_parent(self, idx); self.get_edges()[idx] = child; } } @@ -622,13 +642,29 @@ const BTree = struct { if (start <= self.leaf.len) { for (self.get_edges()[start..], start..) |edg, i| { if (edg) |edge| { - edge.as_leaf().parent = .{ .parent = self, .idx = @intCast(u16, i) }; + edge.as_leaf().set_parent(self, @intCast(u16, i - count)); self.edges[i - count] = edge; } } } } + fn shift_edges_right(self: *Node, start: u16, count: u16) void { + std.debug.assert(self.leaf.len - start >= count - 1); + if (start <= self.leaf.len) { + // shift edges right + for (count..(self.leaf.len + 1 - start)) |i| { + const idx = self.leaf.len - 1 - count - i; + if (self.get_edge(@intCast(u16, idx))) |edge| { + edge.as_leaf().set_parent(self, @intCast(u16, idx + count)); + self.edges[idx + count] = edge; + // set moved eges to null + self.edges[idx] = null; + } + } + } + } + fn dbg(self: *Node) void { const values = self.leaf.get_values(); const edges = self.get_edges()[0..values.len]; @@ -695,6 +731,14 @@ const BTree = struct { self.ally.destroy(self); } + fn set_parent(self: *Leaf, parent: *Node, idx: u16) void { + if (idx < 0 or idx > parent.leaf.len or idx > CAPACITY) { + std.debug.print("{}", .{idx}); + } + + self.parent = .{ .parent = parent, .idx = idx }; + } + const SplitResult = struct { // attached, old node that may be modified left: *Leaf, @@ -762,7 +806,20 @@ const BTree = struct { fn rebalance(self: *Leaf) void { if (self.needs_rebalance()) { if (self.parent) |parent| { - const left_sib = Optional.into_option(parent.parent.get_edge(parent.idx - 1)); + const nonnull = blk: { + if (parent.idx == 0) { + break :blk Optional.Option(u16).none(); + } else { + break :blk Optional.Option(u16).some(parent.idx); + } + }; + const left_sib = nonnull.join_with(struct { + node: *Node, + fn call(a: @This(), idx: u16) Optional.Option(NodeOrLeaf) { + return Optional.into_option(a.node.get_edge(idx - 1)); + } + }{ .node = parent.parent }); + const right_sib = Optional.into_option(parent.parent.get_edge(parent.idx + 1)); const Side = enum { Left, Right }; @@ -795,10 +852,12 @@ const BTree = struct { }; const offset = @intCast(i32, sibling.node.as_leaf().parent.?.idx) - @intCast(i32, parent.idx); + std.debug.print("{any} {} {} {}\n", .{ sibling.node.as_leaf().parent, sibling.node.as_leaf().parent.?.idx, parent.idx, offset }); if (sibling.node.as_leaf().len > MIN_AFTER_SPLIT) { // rotate - const index = @intCast(usize, parent.idx + offset); + const index = @intCast(usize, @max(0, parent.idx + offset)); + std.debug.print("{} {} {}\n", .{ parent.idx, offset, index }); NodeOrLeaf.from_leaf(self).push_value(parent.parent.leaf.values[index]); parent.parent.leaf.values[index] = blk: { @@ -812,6 +871,19 @@ const BTree = struct { break :blk value; } }; + + // TODO: rotate edges + switch (NodeOrLeaf.from_leaf(self)) { + .internal => |iself| { + const isibling = sibling.node.internal; + iself.shift_edges_right(0, 1); + iself.edges[0] = isibling.edges[isibling.leaf.len]; + isibling.edges[isibling.leaf.len] = null; + + iself.edges[0].?.as_leaf().set_parent(iself, 0); + }, + .leaf => {}, + } } else { // merge const left = switch (sibling.side) { @@ -835,6 +907,23 @@ const BTree = struct { // copy values from right to left std.mem.copy(u32, left.values[left.len..], right.values[0..right.len]); + + // TODO: copy edges from right to left aswell + switch (NodeOrLeaf.from_leaf(left)) { + .internal => |ileft| { + const iright = @ptrCast(*Node, right); + + std.mem.copy(?NodeOrLeaf, ileft.edges[0..], iright.edges[0..]); + + for (ileft.get_edges(), 0..) |e, i| { + if (e) |ee| { + ee.as_leaf().set_parent(ileft, @intCast(u16, i)); + } + } + }, + .leaf => {}, + } + left.len = left.len + right.len; parent.parent.as_leaf().rebalance(); @@ -857,7 +946,7 @@ const BTree = struct { for (root.get_edges(), 0..) |e, i| { if (e) |ee| { - ee.as_leaf().parent = .{ .parent = root, .idx = @intCast(u16, i) }; + ee.as_leaf().set_parent(root, @intCast(u16, i)); } } }, @@ -880,6 +969,7 @@ const BTree = struct { } fn remove_key(self: *Leaf, key: u32) void { + // TODO: remove keys from internal nodes // safety: key must be contained by .values var n: u16 = 0; for (self.get_values(), 0..) |val, i| { @@ -1040,8 +1130,24 @@ test "btree rand insert" { } test "ref removal" { - std.testing.refAllDeclsRecursive(BTree.Leaf); - std.testing.refAllDeclsRecursive(BTree.NodeOrLeaf); + std.debug.print("remove keys\n", .{}); + + var tree = BTree.create(std.testing.allocator); + defer tree.destroy(); + + for (0..20) |i| { + _ = try tree.insert(@intCast(u32, i)); + } + + tree.dbg(); + std.debug.print("\n", .{}); + + for (0..15) |i| { + tree.remove(@intCast(u32, i)); + } + + tree.dbg(); + std.debug.print("\n", .{}); } test "btree rebalance" {