Eugene Rossokha
369 lines
11 KiB
Zig
369 lines
11 KiB
Zig
const std = @import("std");
|
|
|
|
const Allocator = std.mem.Allocator;
|
|
pub const Order = std.math.Order;
|
|
|
|
pub fn CompareFn(comptime N: type) type {
|
|
return fn (*const N, *const N) Order;
|
|
}
|
|
|
|
pub fn SearchFn(comptime N: type, comptime C: type) type {
|
|
return fn (*const N, C) Order;
|
|
}
|
|
|
|
pub fn BTree(comptime N: type, comptime compare_fn: CompareFn(N), comptime deinit_fn: ?fn (*N) void) type {
|
|
return struct {
|
|
gpa: Allocator,
|
|
root: ?*Node = null,
|
|
|
|
pub const Error = error{ already_exists, does_not_exist } || Allocator.Error;
|
|
|
|
pub fn WalkFn(comptime C: type) type {
|
|
return fn (*const Node, C) void;
|
|
}
|
|
|
|
pub const Iterator = struct {
|
|
current: ?*Node,
|
|
|
|
pub fn next(self: *Iterator) ?*const Node {
|
|
while (self.current) |n| {
|
|
const v = n;
|
|
|
|
if (n.right) |r| {
|
|
// Emit
|
|
self.current = Node.leftmost(r);
|
|
} else {
|
|
var nn = n;
|
|
while (nn.parent) |p| {
|
|
if (nn == p.right) {
|
|
nn = p;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
self.current = nn.parent;
|
|
}
|
|
|
|
return v;
|
|
}
|
|
|
|
return null;
|
|
}
|
|
};
|
|
|
|
pub const Node = struct {
|
|
key: N,
|
|
parent: ?*Node = null,
|
|
left: ?*Node = null,
|
|
right: ?*Node = null,
|
|
|
|
fn init(a: Allocator, key: N) Error!*Node {
|
|
const node = try a.create(Node);
|
|
node.* = .{
|
|
.key = key,
|
|
};
|
|
return node;
|
|
}
|
|
|
|
fn deinit(node: ?*Node, a: Allocator) void {
|
|
if (node) |n| {
|
|
if (comptime deinit_fn) |f| {
|
|
f(&n.key);
|
|
}
|
|
|
|
Node.deinit(n.left, a);
|
|
Node.deinit(n.right, a);
|
|
|
|
// Free node itself
|
|
a.destroy(n);
|
|
}
|
|
}
|
|
|
|
fn insert(node: ?*Node, a: Allocator, key: N) Error!struct { *Node, *Node } {
|
|
if (node) |n| {
|
|
const ord = compare_fn(&n.key, &key);
|
|
var inserted: *Node = undefined;
|
|
switch (ord) {
|
|
.lt => {
|
|
const child, inserted = try Node.insert(n.right, a, key);
|
|
child.parent = n;
|
|
n.right = child;
|
|
},
|
|
.gt => {
|
|
const child, inserted = try Node.insert(n.left, a, key);
|
|
child.parent = n;
|
|
n.left = child;
|
|
},
|
|
.eq => return error.already_exists,
|
|
}
|
|
return .{ n, inserted };
|
|
} else {
|
|
const n = try Node.init(a, key);
|
|
return .{ n, n };
|
|
}
|
|
}
|
|
|
|
fn remove_node(node: *Node, a: Allocator, destroy: bool) ?*Node {
|
|
if (node.left == null) {
|
|
// Only right/none
|
|
const tmp = node.right;
|
|
if (tmp) |t| {
|
|
t.parent = node.parent;
|
|
}
|
|
// Destroy the node
|
|
if (comptime deinit_fn) |f| {
|
|
if (destroy) {
|
|
f(&node.key);
|
|
}
|
|
}
|
|
a.destroy(node);
|
|
return tmp;
|
|
}
|
|
if (node.right == null) {
|
|
// Only left/none
|
|
const tmp = node.left;
|
|
if (tmp) |t| {
|
|
t.parent = node.parent;
|
|
}
|
|
// Destroy the node
|
|
if (comptime deinit_fn) |f| {
|
|
if (destroy) {
|
|
f(&node.key);
|
|
}
|
|
}
|
|
a.destroy(node);
|
|
return tmp;
|
|
}
|
|
|
|
// Both
|
|
var successor = node.right;
|
|
while (successor) |succ| {
|
|
if (succ.left) |l| {
|
|
successor = l;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
if (successor) |succ| {
|
|
node.key = succ.key;
|
|
node.right = Node.remove(node.right, a, succ.key) catch unreachable;
|
|
}
|
|
return node;
|
|
}
|
|
|
|
fn remove(node: ?*Node, a: Allocator, key: N) Error!?*Node {
|
|
if (node) |n| {
|
|
const ord = compare_fn(&n.key, &key);
|
|
switch (ord) {
|
|
.lt => n.right = try Node.remove(n.right, a, key),
|
|
.gt => n.left = try Node.remove(n.left, a, key),
|
|
.eq => return Node.remove_node(n, a, true),
|
|
}
|
|
return node;
|
|
} else {
|
|
return error.does_not_exist;
|
|
}
|
|
}
|
|
|
|
fn walk(node: ?*Node, ctx: anytype, walk_fn: WalkFn(@TypeOf(ctx))) void {
|
|
if (node) |n| {
|
|
Node.walk(n.left, ctx, walk_fn);
|
|
walk_fn(n, ctx);
|
|
Node.walk(n.right, ctx, walk_fn);
|
|
}
|
|
}
|
|
|
|
fn leftmost(node: ?*Node) ?*Node {
|
|
var n = node;
|
|
while (n) |nn| {
|
|
if (nn.left == null) {
|
|
break;
|
|
}
|
|
n = nn.left;
|
|
}
|
|
return n;
|
|
}
|
|
};
|
|
|
|
pub fn init(a: std.mem.Allocator) @This() {
|
|
return .{ .gpa = a };
|
|
}
|
|
|
|
pub fn deinit(self: *@This()) void {
|
|
Node.deinit(self.root, self.gpa);
|
|
}
|
|
|
|
pub fn iterator(self: *@This()) Iterator {
|
|
return .{ .current = Node.leftmost(self.root) };
|
|
}
|
|
|
|
pub fn insert(self: *@This(), key: N) Error!*Node {
|
|
self.root, const inserted = try Node.insert(self.root, self.gpa, key);
|
|
return inserted;
|
|
}
|
|
|
|
pub fn remove(self: *@This(), key: N) Error!void {
|
|
self.root = try Node.remove(self.root, self.gpa, key);
|
|
}
|
|
|
|
pub fn remove_node(self: *@This(), node: *Node, destroy: bool) Error!void {
|
|
if (node.parent) |p| {
|
|
// Non-root node
|
|
const np = Node.remove_node(node, self.gpa, destroy);
|
|
if (np) |npp| {
|
|
npp.parent = p;
|
|
}
|
|
if (node == p.right) {
|
|
p.right = np;
|
|
} else {
|
|
p.left = np;
|
|
}
|
|
} else {
|
|
// Root node
|
|
const np = Node.remove_node(node, self.gpa, destroy);
|
|
if (np) |npp| {
|
|
npp.parent = null;
|
|
}
|
|
self.root = np;
|
|
}
|
|
}
|
|
|
|
pub fn lookup(self: *const @This(), key: N) ?*Node {
|
|
const search_fn = struct {
|
|
fn call(n: *const N, cx: N) Order {
|
|
return compare_fn(n, &cx);
|
|
}
|
|
}.call;
|
|
|
|
return self.search(key, search_fn);
|
|
}
|
|
|
|
pub fn search(
|
|
self: *const @This(),
|
|
ctx: anytype,
|
|
search_fn: SearchFn(N, @TypeOf(ctx)),
|
|
) ?*Node {
|
|
var node = self.root;
|
|
while (node) |n| {
|
|
const ord = search_fn(&n.key, ctx);
|
|
switch (ord) {
|
|
.gt => node = n.left,
|
|
.eq => return n,
|
|
.lt => node = n.right,
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
pub fn walk(self: *@This(), ctx: anytype, walk_fn: WalkFn(@TypeOf(ctx))) void {
|
|
Node.walk(self.root, ctx, walk_fn);
|
|
}
|
|
};
|
|
}
|
|
|
|
test "BTree insertion/removal" {
|
|
const int_compare_fn = struct {
|
|
fn call(a: *const u32, b: *const u32) Order {
|
|
if (a.* > b.*) {
|
|
return .gt;
|
|
} else if (a.* == b.*) {
|
|
return .eq;
|
|
} else {
|
|
return .lt;
|
|
}
|
|
}
|
|
}.call;
|
|
const Tree = BTree(u32, int_compare_fn, null);
|
|
var tree = Tree.init(std.testing.allocator);
|
|
defer tree.deinit();
|
|
|
|
for (50..100) |i| {
|
|
_ = try tree.insert(@truncate(i));
|
|
}
|
|
for (1..50) |i| {
|
|
_ = try tree.insert(@truncate(i));
|
|
}
|
|
|
|
for (1..100) |i| {
|
|
const k = @as(u32, @truncate(i));
|
|
try std.testing.expectEqual(k, tree.lookup(k).?.key);
|
|
}
|
|
|
|
for (1..100) |i| {
|
|
const k = 100 - @as(u32, @truncate(i));
|
|
if (i % 2 == 0) {
|
|
try tree.remove(k);
|
|
}
|
|
}
|
|
|
|
for (1..100) |i| {
|
|
const k = @as(u32, @truncate(i));
|
|
if (i % 2 == 0) {
|
|
try std.testing.expectEqual(null, tree.lookup(k));
|
|
} else {
|
|
try std.testing.expectEqual(k, tree.lookup(k).?.key);
|
|
}
|
|
}
|
|
}
|
|
|
|
test "BTree removal by node" {
|
|
const int_compare_fn = struct {
|
|
fn call(a: *const u32, b: *const u32) Order {
|
|
if (a.* > b.*) {
|
|
return .gt;
|
|
} else if (a.* == b.*) {
|
|
return .eq;
|
|
} else {
|
|
return .lt;
|
|
}
|
|
}
|
|
}.call;
|
|
const Tree = BTree(u32, int_compare_fn, null);
|
|
var tree = Tree.init(std.testing.allocator);
|
|
defer tree.deinit();
|
|
|
|
_ = try tree.insert(10);
|
|
_ = try tree.insert(11);
|
|
_ = try tree.insert(12);
|
|
|
|
{
|
|
const n = tree.lookup(10).?;
|
|
try tree.remove_node(n, true);
|
|
}
|
|
|
|
try std.testing.expectEqual(null, tree.lookup(10));
|
|
try std.testing.expectEqual(12, tree.lookup(12).?.key);
|
|
try std.testing.expectEqual(11, tree.lookup(11).?.key);
|
|
}
|
|
|
|
test "BTree iterator" {
|
|
const int_compare_fn = struct {
|
|
fn call(a: *const u32, b: *const u32) Order {
|
|
if (a.* > b.*) {
|
|
return .gt;
|
|
} else if (a.* == b.*) {
|
|
return .eq;
|
|
} else {
|
|
return .lt;
|
|
}
|
|
}
|
|
}.call;
|
|
const Tree = BTree(u32, int_compare_fn, null);
|
|
var tree = Tree.init(std.testing.allocator);
|
|
defer tree.deinit();
|
|
|
|
for (50..100) |i| {
|
|
_ = try tree.insert(@truncate(i));
|
|
}
|
|
for (1..50) |i| {
|
|
_ = try tree.insert(@truncate(i));
|
|
}
|
|
|
|
var it = tree.iterator();
|
|
for (1..100) |i| {
|
|
const n = it.next().?;
|
|
try std.testing.expectEqual(i, n.key);
|
|
}
|
|
try std.testing.expectEqual(null, it.next());
|
|
}
|