read all from socket

.gitignore

@@ -1,4 +1,4 @@
-pizzakv
+pizzakv*
 main
 tools/
 .db

index.zig

@@ -1,146 +1,238 @@
 const std = @import("std");
 const storage = @import("storage.zig");
 
-var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
-const allocator = arena.allocator();
-
-const MAX_CHILDREN = 256;
+var tree_arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
+const tree_allocator = tree_arena.allocator();
+
+const temp_allocator = std.heap.c_allocator;
+
+const RadixNode = struct {
+    edge: []const u8,
+    children: std.StringHashMap(*RadixNode),
+    is_terminal: bool,
+
+    fn init(edge: []const u8) *RadixNode {
+        const node = tree_allocator.create(RadixNode) catch unreachable;
+        node.* = .{
+            .edge = tree_allocator.dupe(u8, edge) catch unreachable,
+            .children = std.StringHashMap(*RadixNode).init(tree_allocator),
+            .is_terminal = false,
+        };
+        return node;
+    }
 
-const TrieNode = struct {
-    children: [MAX_CHILDREN]?*TrieNode,
-    eof: bool,
-    char: u8,
+    fn deinit(self: *RadixNode) void {
+        var it = self.children.iterator();
+        while (it.next()) |entry| {
+            entry.value_ptr.*.deinit();
+        }
+        self.children.deinit();
+        tree_allocator.free(self.edge);
+        tree_allocator.destroy(self);
+    }
 };
 
-var root: TrieNode = .{
-    .children = [_]?*TrieNode{null} ** MAX_CHILDREN,
-    .eof = false,
-    .char = 0,
-};
+var root: *RadixNode = undefined;
+var root_initialized = false;
 
-fn getChild(node: *TrieNode, c: u8) ?*TrieNode {
-    return node.children[c];
+fn ensureRoot() void {
+    if (!root_initialized) {
+        root = RadixNode.init("");
+        root_initialized = true;
+    }
 }
 
-fn addChild(node: *TrieNode, c: u8) *TrieNode {
-    const newNode = allocator.create(TrieNode) catch unreachable;
-    newNode.* = .{
-        .children = [_]?*TrieNode{null} ** MAX_CHILDREN,
-        .eof = false,
-        .char = c,
-    };
-    node.children[c] = newNode;
-    return newNode;
+fn commonPrefixLen(a: []const u8, b: []const u8) usize {
+    var i: usize = 0;
+    while (i < a.len and i < b.len and a[i] == b[i]) {
+        i += 1;
+    }
+    return i;
 }
 
 pub fn insert(key: []const u8) void {
-    var current = &root;
-    for (key) |c| {
-        var child = getChild(current, c);
-        if (child == null) {
-            child = addChild(current, c);
+    ensureRoot();
+    if (key.len == 0) return;
+
+    var node = root;
+    var remaining = key;
+
+    while (remaining.len > 0) {
+        var found = false;
+
+        var it = node.children.iterator();
+        while (it.next()) |entry| {
+            const child = entry.value_ptr.*;
+            const prefix_len = commonPrefixLen(child.edge, remaining);
+
+            if (prefix_len > 0) {
+                found = true;
+
+                if (prefix_len == child.edge.len) {
+                    if (prefix_len == remaining.len) {
+                        child.is_terminal = true;
+                        return;
+                    }
+                    remaining = remaining[prefix_len..];
+                    node = child;
+                    break;
+                } else {
+                    const old_edge = child.edge;
+                    const common = old_edge[0..prefix_len];
+                    const child_suffix = old_edge[prefix_len..];
+                    const key_suffix = remaining[prefix_len..];
+
+                    const intermediate = RadixNode.init(common);
+
+                    tree_allocator.free(child.edge);
+                    child.edge = tree_allocator.dupe(u8, child_suffix) catch unreachable;
+
+                    intermediate.children.put(child_suffix, child) catch unreachable;
+
+                    _ = node.children.remove(old_edge);
+                    node.children.put(common, intermediate) catch unreachable;
+
+                    if (key_suffix.len == 0) {
+                        intermediate.is_terminal = true;
+                        return;
+                    } else {
+                        const new_child = RadixNode.init(key_suffix);
+                        new_child.is_terminal = true;
+                        intermediate.children.put(key_suffix, new_child) catch unreachable;
+                        return;
+                    }
+                }
+            }
+        }
+
+        if (!found) {
+            const new_child = RadixNode.init(remaining);
+            new_child.is_terminal = true;
+            node.children.put(remaining, new_child) catch unreachable;
+            return;
         }
-        current = child.?;
     }
-    current.eof = true;
+
+    node.is_terminal = true;
 }
 
 pub fn delete(key: []const u8) void {
-    var current = &root;
-    for (key) |c| {
-        const child = getChild(current, c);
-        if (child == null) {
-            return;
-        }
-        current = child.?;
+    ensureRoot();
+    if (key.len == 0) return;
+
+    const node = findNode(root, key);
+    if (node) |n| {
+        n.is_terminal = false;
     }
-    current.eof = false;
 }
 
-pub fn searchByPrefix(prefix: []const u8) ?*TrieNode {
-    var current = &root;
-    for (prefix) |c| {
-        const child = getChild(current, c);
-        if (child == null) {
-            return null;
+fn findNode(node: *RadixNode, key: []const u8) ?*RadixNode {
+    if (key.len == 0) return node;
+
+    var current = node;
+    var remaining = key;
+
+    while (remaining.len > 0) {
+        var found = false;
+
+        var it = current.children.iterator();
+        while (it.next()) |entry| {
+            const child = entry.value_ptr.*;
+            const prefix_len = commonPrefixLen(child.edge, remaining);
+
+            if (prefix_len > 0) {
+                if (prefix_len < child.edge.len) {
+                    return null;
+                }
+
+                if (prefix_len == remaining.len) {
+                    return child;
+                }
+
+                remaining = remaining[prefix_len..];
+                current = child;
+                found = true;
+                break;
+            }
         }
-        current = child.?;
+
+        if (!found) return null;
     }
 
     return current;
 }
 
-fn countKeys(node: *TrieNode) usize {
+pub fn searchByPrefix(prefix: []const u8) ?*RadixNode {
+    ensureRoot();
+    if (prefix.len == 0) return root;
+    return findNode(root, prefix);
+}
+
+fn countKeys(node: *RadixNode) usize {
     var count: usize = 0;
-    if (node.eof) {
+    if (node.is_terminal) {
         count += 1;
     }
 
-    for (0..MAX_CHILDREN) |i| {
-        const child = node.children[i];
-        if (child != null) {
-            count += countKeys(child.?);
-        }
+    var it = node.children.iterator();
+    while (it.next()) |entry| {
+        count += countKeys(entry.value_ptr.*);
     }
     return count;
 }
 
-fn collectKeys(node: *TrieNode, prefix: []const u8, keys: [][]const u8, index: *usize) void {
-    if (node.eof) {
-        const key = allocator.alloc(u8, prefix.len) catch unreachable;
-        @memcpy(key, prefix);
-        keys[index.*] = key;
-        index.* += 1;
-    }
+const MAX_KEYS_RETURN = 10000;
 
-    for (0..MAX_CHILDREN) |i| {
-        const child = node.children[i];
-        if (child != null) {
-            const childChar = child.?.char;
-            var newPrefix = allocator.alloc(u8, prefix.len + 1) catch unreachable;
-            @memcpy(newPrefix[0..prefix.len], prefix);
-            newPrefix[prefix.len] = childChar;
-            collectKeys(child.?, newPrefix, keys, index);
-            allocator.free(newPrefix);
-        }
+fn collectKeys(node: *RadixNode, prefix: []const u8, keys: *std.ArrayListUnmanaged([]const u8), max_keys: usize) void {
+    if (keys.items.len >= max_keys) return;
+
+    if (node.is_terminal) {
+        const key = temp_allocator.dupe(u8, prefix) catch return;
+        keys.append(temp_allocator, key) catch return;
     }
-}
 
-pub fn getKeysFromNode(node: *TrieNode, prefix: []const u8) [][]const u8 {
-    const keyCount = countKeys(node);
-    if (keyCount == 0) {
-        return &[_][]const u8{};
+    var it = node.children.iterator();
+    while (it.next()) |entry| {
+        if (keys.items.len >= max_keys) break;
+        const child = entry.value_ptr.*;
+        const new_prefix = std.mem.concat(temp_allocator, u8, &[_][]const u8{ prefix, child.edge }) catch return;
+        collectKeys(child, new_prefix, keys, max_keys);
+        temp_allocator.free(new_prefix);
     }
+}
 
-    const keys = allocator.alloc([]const u8, keyCount) catch unreachable;
-    var index: usize = 0;
-    collectKeys(node, prefix, keys, &index);
-    return keys;
+pub fn getKeysFromNode(node: *RadixNode, prefix: []const u8) [][]const u8 {
+    var keys_list = std.ArrayListUnmanaged([]const u8){};
+    collectKeys(node, prefix, &keys_list, MAX_KEYS_RETURN);
+    return keys_list.toOwnedSlice(temp_allocator) catch &[_][]const u8{};
 }
 
 pub fn getKeysByPrefix(prefix: []const u8) []const u8 {
+    ensureRoot();
     const node = searchByPrefix(prefix) orelse return "";
     const keys = getKeysFromNode(node, prefix);
     if (keys.len == 0) return "";
-    return std.mem.join(allocator, "\r", keys) catch unreachable;
+    return std.mem.join(temp_allocator, "\r", keys) catch "";
 }
 
 pub fn getValuesByPrefix(prefix: []const u8) []const u8 {
+    ensureRoot();
     const node = searchByPrefix(prefix) orelse return "";
     const keys = getKeysFromNode(node, prefix);
     if (keys.len == 0) return "";
 
-    const values = allocator.alloc([]const u8, keys.len) catch unreachable;
+    const values = temp_allocator.alloc([]const u8, keys.len) catch return "";
     for (keys, 0..) |key, i| {
         const value = storage.read(key) orelse "";
         values[i] = value;
     }
 
-    return std.mem.join(allocator, "\r", values) catch unreachable;
+    return std.mem.join(temp_allocator, "\r", values) catch "";
 }
 
 pub fn getAllKeys() []const u8 {
-    const keys = getKeysFromNode(&root, &[_]u8{});
+    ensureRoot();
+    const keys = getKeysFromNode(root, &[_]u8{});
     if (keys.len == 0) return "";
-    return std.mem.join(allocator, "\r", keys) catch unreachable;
+    return std.mem.join(temp_allocator, "\r", keys) catch "";
 }

socket.zig

@@ -18,18 +18,22 @@ pub fn init(port: u16) !posix.socket_t {
 }
 
 pub fn readUntilCR(conn: posix.socket_t, buf: []u8) !usize {
-    var pos: usize = 0;
-    while (pos < buf.len) {
-        const n = try posix.read(conn, buf[pos .. pos + 1]);
+    var total: usize = 0;
+
+    while (total < buf.len) {
+        const n = try posix.read(conn, buf[total..]);
         if (n == 0) {
-            return pos;
+            return if (total > 0) total else error.ConnectionClosed;
         }
-        if (buf[pos] == '\r') {
-            return pos + 1;
+
+        if (std.mem.indexOfScalar(u8, buf[total .. total + n], '\r')) |offset| {
+            return total + offset;
         }
-        pos += n;
+
+        total += n;
     }
-    return pos;
+
+    return total;
 }
 
 pub fn readUntilNewLine(conn: posix.socket_t, buf: []u8) !usize {