janis/zig-btree
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

removal works better but still somehow broken somewhere

Browse source
This commit is contained in:
Janis 2023-03-20 01:15:18 +01:00
parent d3670bcf69
commit 0ee791db3c
1 changed files with 131 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

150
src/main.zig
View file

@ -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) {
const myleaf = blk: {
if (right_value.leaf.len > left_value.leaf.len) {
// use right leaf
right_value;
break :blk right_value;
} else {
// use left leaf
left_value;
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" {