use MemoryPool instead of allocator, fix examples

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});
 }

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) {

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);

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,41 +42,39 @@ 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);
             }
             return null;
         }
 
         /// 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;
+            }
+        };
+    };
 }

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) {

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);