diff --git a/examples/basic_usage.zig b/examples/basic_usage.zig index 7628306..aace66a 100644 --- a/examples/basic_usage.zig +++ b/examples/basic_usage.zig @@ -1,14 +1,20 @@ const std = @import("std"); -const bpt = @import("bplustree").BPlusTree(i32, 4); +const bpt = @import("bplustree").BPlusTree(i32, i32, 4); pub fn main() !void { - var tree = bpt.init(&std.heap.page_allocator); + var dbga: std.heap.DebugAllocator(.{}) = .init; + defer _ = dbga.deinit(); + + var tree = bpt.init(dbga.allocator()); defer tree.deinit(); + try tree.insert(10, 100); try tree.insert(20, 200); try tree.insert(5, 50); + const found = tree.search(10); std.debug.print("Found 10: {any}\n", .{found}); + const not_found = tree.search(99); std.debug.print("Not found 99: {any}\n", .{not_found}); } diff --git a/examples/iteration.zig b/examples/iteration.zig index 7a92915..8f3fc45 100644 --- a/examples/iteration.zig +++ b/examples/iteration.zig @@ -1,13 +1,18 @@ const std = @import("std"); -const bpt = @import("bplustree").BPlusTree(i32, 4); +const bpt = @import("bplustree").BPlusTree(i32, i32, 4); pub fn main() !void { - var tree = bpt.init(&std.heap.page_allocator); + var dbga: std.heap.DebugAllocator(.{}) = .init; + defer _ = dbga.deinit(); + + var tree = bpt.init(dbga.allocator()); defer tree.deinit(); + try tree.insert(1, 10); try tree.insert(2, 20); try tree.insert(3, 30); try tree.insert(4, 40); + // This tree does not yet have an iterator, but you can search for all keys: var k: i32 = 1; while (k < 5) : (k += 1) { diff --git a/examples/self_balancing.zig b/examples/self_balancing.zig index 3a488ae..d12363c 100644 --- a/examples/self_balancing.zig +++ b/examples/self_balancing.zig @@ -1,12 +1,15 @@ const std = @import("std"); -const bpt = @import("bplustree").BPlusTree(i32, 4); +const bpt = @import("bplustree").BPlusTree(i32, i32, 4); pub fn main() !void { - var tree = bpt.init(&std.heap.page_allocator); + var dbga: std.heap.DebugAllocator(.{}) = .init; + defer _ = dbga.deinit(); + + var tree = bpt.init(dbga.allocator()); defer tree.deinit(); // Insert enough keys to force multiple splits and balancing - const N = 100; + const N = 1024 * 1024; var k: i32 = 0; while (k < N) : (k += 1) { try tree.insert(k, k * 10); diff --git a/src/bplustree.zig b/src/bplustree.zig index 627562e..fcb6e0d 100644 --- a/src/bplustree.zig +++ b/src/bplustree.zig @@ -1,56 +1,35 @@ const std = @import("std"); -/// Returns a B+ tree node type for the given key/value type and degree. -pub fn Node(comptime K: type, comptime V: type, comptime DEGREE: usize) type { +fn NodeType(comptime K: type, comptime V: type, comptime DEGREE: usize) type { return struct { is_leaf: bool, keys: [2 * DEGREE - 1]K, - children: [2 * DEGREE]*Node(K, V, DEGREE), + children: [2 * DEGREE]*NodeType(K, V, DEGREE), values: [2 * DEGREE - 1]?V, n: usize, - next: ?*Node(K, V, DEGREE), + next: ?*NodeType(K, V, DEGREE), }; } /// Returns a robust, generic B+ tree type for the given key/value type and degree. /// Usage: var tree = BPlusTree(i32, 4).init(allocator); pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) type { - const Allocator = std.mem.Allocator; - - const Error = error{ - OutOfMemory, - NotFound, - DuplicateKey, - }; + return struct { + pub const Error = error{ + OutOfMemory, + NotFound, + DuplicateKey, + }; - const NodePtr = *Node(K, V, DEGREE); + pub const Node = NodeType(K, V, DEGREE); + const NodePtr = *Node; + const Self = @This(); - const Tree = struct { - /// Iterator for in-order traversal of the B+ tree. - pub const Iterator = struct { - current: ?NodePtr, - idx: usize, - - /// Returns the next key-value pair, or null if done. - pub fn next(self: *@This()) ?struct { key: K, value: V } { - if (self.current) |node| { - if (self.idx < node.n) { - const result = .{ .key = node.keys[self.idx], .value = node.values[self.idx].? }; - self.idx += 1; - return result; - } else if (node.next) |next_node| { - self.current = next_node; - self.idx = 0; - return self.next(); - } - } - return null; - } - }; root: ?NodePtr, - allocator: *const Allocator, + pool: std.heap.MemoryPool(Node), + /// Returns an iterator for in-order traversal. - pub fn iter(self: *@This()) Iterator { + pub fn iter(self: *Self) Iterator { var node = self.root; // Find leftmost leaf while (node != null and !node.?.is_leaf) { @@ -63,33 +42,31 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Initialize a new B+ tree with the given allocator. - pub fn init(allocator: *const Allocator) @This() { - return @This(){ + pub fn init(gpa: std.mem.Allocator) Self { + return Self{ .root = null, - .allocator = allocator, + .pool = .init(gpa), }; } /// Free all memory used by the tree and its nodes. - pub fn deinit(self: *@This()) void { - if (self.root) |r| { - self.freeNode(r); - } + pub fn deinit(self: *Self) void { + self.pool.deinit(); + self.* = undefined; } - pub const ErrorSet = Error; /// Recursively free a node and its children. - fn freeNode(self: *@This(), node: NodePtr) void { + fn freeNode(self: *Self, node: NodePtr) void { if (!node.is_leaf) { for (node.children[0 .. node.n + 1]) |child| { self.freeNode(child); } } - self.allocator.destroy(node); + self.pool.destroy(node); } /// Search for a key in the tree. Returns the value if found, else null. - pub fn search(self: *@This(), key: K) ?V { + pub fn search(self: *Self, key: K) ?V { if (self.root) |r| { return self.searchNode(r, key); } @@ -97,7 +74,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Internal recursive search helper. - fn searchNode(self: *@This(), node: NodePtr, key: K) ?V { + fn searchNode(self: *Self, node: NodePtr, key: K) ?V { var i: usize = 0; while (i < node.n and key > node.keys[i]) : (i += 1) {} if (node.is_leaf) { @@ -116,7 +93,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Insert a key-value pair into the tree. Returns error on duplicate key or OOM. - pub fn insert(self: *@This(), key: K, value: V) Error!void { + pub fn insert(self: *Self, key: K, value: V) Error!void { if (self.root == null) { self.root = try self.createNode(true); self.root.?.keys[0] = key; @@ -134,7 +111,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Insert a key-value pair into a node that is not full. - fn insertNonFull(self: *@This(), node: NodePtr, key: K, value: V) Error!void { + fn insertNonFull(self: *Self, node: NodePtr, key: K, value: V) Error!void { var i = node.n; if (node.is_leaf) { while (i > 0 and key < node.keys[i - 1]) : (i -= 1) { @@ -158,7 +135,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ /// Split a full child node and update the parent. /// Parent must never be a leaf in a B+ tree. - fn splitChild(self: *@This(), parent: NodePtr, i: usize, y: NodePtr) Error!void { + fn splitChild(self: *Self, parent: NodePtr, i: usize, y: NodePtr) Error!void { std.debug.assert(!parent.is_leaf); var z = try self.createNode(y.is_leaf); var j: usize = 0; @@ -209,9 +186,9 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Allocate and initialize a new node (leaf or internal). - fn createNode(self: *@This(), is_leaf: bool) !NodePtr { - const node = try self.allocator.create(Node(K, V, DEGREE)); - node.* = Node(K, V, DEGREE){ + fn createNode(self: *Self, is_leaf: bool) !NodePtr { + const node = try self.pool.create(); + node.* = Node{ .is_leaf = is_leaf, .keys = undefined, .children = undefined, @@ -223,7 +200,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Remove a key from the tree. Returns error if not found. - pub fn remove(self: *@This(), key: K) Error!void { + pub fn remove(self: *Self, key: K) Error!void { if (self.root == null) return Error.NotFound; try self.removeNode(self.root.?, key); // If root is empty and not a leaf, collapse tree height @@ -232,13 +209,13 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } // If root is empty and is a leaf, tree is now empty if (self.root.?.n == 0 and self.root.?.is_leaf) { - self.allocator.destroy(self.root.?); + self.freeNode(self.root.?); self.root = null; } } /// Internal recursive remove helper. - fn removeNode(self: *@This(), node: NodePtr, key: K) Error!void { + fn removeNode(self: *Self, node: NodePtr, key: K) Error!void { var i: usize = 0; while (i < node.n and key > node.keys[i]) : (i += 1) {} if (node.is_leaf) { @@ -283,7 +260,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Borrow a key from the previous sibling - fn borrowFromPrev(self: *@This(), idx: usize, _parent: NodePtr) void { + fn borrowFromPrev(self: *Self, idx: usize, _parent: NodePtr) void { _ = self; // autofix const child = _parent.children[idx]; const sibling = _parent.children[idx - 1]; @@ -308,7 +285,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Borrow a key from the next sibling - fn borrowFromNext(self: *@This(), idx: usize, _parent: NodePtr) void { + fn borrowFromNext(self: *Self, idx: usize, _parent: NodePtr) void { _ = self; // autofix const child = _parent.children[idx]; const sibling = _parent.children[idx + 1]; @@ -334,7 +311,7 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ } /// Merge child at idx with its right sibling - fn merge(self: *@This(), node: NodePtr, idx: usize) void { + fn merge(self: *Self, node: NodePtr, idx: usize) void { const child = node.children[idx]; const sibling = node.children[idx + 1]; // For internal nodes, bring down separator key @@ -363,9 +340,33 @@ pub fn BPlusTree(comptime K: type, comptime V: type, comptime DEGREE: usize) typ node.children[k + 1] = node.children[k + 2]; } node.n -= 1; - self.allocator.destroy(sibling); + self.freeNode(sibling); } - }; - return Tree; + /// Iterator for in-order traversal of the B+ tree. + pub const Iterator = struct { + pub const KV = struct { key: K, value: V }; + + current: ?NodePtr, + idx: usize, + + /// Returns the next key-value pair, or null if done. + pub fn next(self: *Iterator) ?KV { + if (self.current) |node| { + if (self.idx < node.n) { + defer self.idx += 1; + return .{ + .key = node.keys[self.idx], + .value = node.values[self.idx].?, + }; + } else if (node.next) |next_node| { + self.current = next_node; + self.idx = 0; + return self.next(); + } + } + return null; + } + }; + }; } diff --git a/tests/bplustree_edge_cases.zig b/tests/bplustree_edge_cases.zig index 5849bcd..c890569 100644 --- a/tests/bplustree_edge_cases.zig +++ b/tests/bplustree_edge_cases.zig @@ -5,21 +5,21 @@ const std = @import("std"); const bpt = @import("bplustree").BPlusTree(i32, i32, 4); test "duplicate key insert fails" { - var tree = bpt.init(&std.testing.allocator); + var tree = bpt.init(std.testing.allocator); defer tree.deinit(); try tree.insert(1, 100); const err = tree.insert(1, 200); - try std.testing.expectError(bpt.ErrorSet.DuplicateKey, err); + try std.testing.expectError(bpt.Error.DuplicateKey, err); } test "search in empty tree returns null" { - var tree = bpt.init(&std.testing.allocator); + var tree = bpt.init(std.testing.allocator); defer tree.deinit(); try std.testing.expect(tree.search(42) == null); } test "insert many and search all" { - var tree = bpt.init(&std.testing.allocator); + var tree = bpt.init(std.testing.allocator); defer tree.deinit(); var i: i32 = 0; while (i < 100) : (i += 1) { diff --git a/tests/bplustree_test.zig b/tests/bplustree_test.zig index 6aa9603..10c2c29 100644 --- a/tests/bplustree_test.zig +++ b/tests/bplustree_test.zig @@ -5,7 +5,7 @@ pub fn main() !void {} // Basic insert/search test test "insert and search" { - var tree = bpt.init(&std.testing.allocator); + var tree = bpt.init(std.testing.allocator); defer tree.deinit(); try tree.insert(1, 100); try tree.insert(2, 200);