Implement index

command.zig

@@ -1,5 +1,6 @@
 const std = @import("std");
 const storage = @import("storage.zig");
+const index = @import("index.zig");
 
 const FAILURE_RESPONSE = "error";
 const SUCCESS_RESPONSE = "success";
@@ -9,6 +10,8 @@ const Command = enum {
     write,
     delete,
     status,
+    keys,
+    reads,
 };
 
 fn parseKeyValue(buf: []const u8) ?[2][]const u8 {
@@ -63,6 +66,13 @@ pub fn parse(msg: []const u8) ?[]const u8 {
 
             return SUCCESS_RESPONSE;
         },
+        .keys => {
+            return index.getAllKeys();
+        },
+        .reads => {
+            const prefix = messageIterator.rest();
+            return index.getValuesByPrefix(prefix);
+        },
         .status => {
             return "well going our operation";
         },

index.zig

@@ -1,32 +1,146 @@
 const std = @import("std");
-const storage = @import("storage");
+const storage = @import("storage.zig");
 
 var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
 const allocator = arena.allocator();
 
-const MAX_CHILDREN = 1024;
+const MAX_CHILDREN = 256;
 
-const RadixNode = struct {
-    children: [MAX_CHILDREN]?*RadixNode,
+const TrieNode = struct {
+    children: [MAX_CHILDREN]?*TrieNode,
     eof: bool,
-    prefix: []const u8,
+    char: u8,
 };
 
-var root: RadixNode = .{
-    .children = [_]?*RadixNode{null} ** MAX_CHILDREN,
+var root: TrieNode = .{
+    .children = [_]?*TrieNode{null} ** MAX_CHILDREN,
     .eof = false,
-    .prefix = &[_]u8{},
+    .char = 0,
 };
 
-fn commonPrefixLen(a: []const u8, b: []const u8) usize {
-    var len: usize = 0;
-    const min_len = if (a.len < b.len) a.len else b.len;
+fn getChild(node: *TrieNode, c: u8) ?*TrieNode {
+    return node.children[c];
+}
+
+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;
+}
+
+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);
+        }
+        current = child.?;
+    }
+    current.eof = 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.?;
+    }
+    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;
+        }
+        current = child.?;
+    }
+
+    return current;
+}
+
+fn countKeys(node: *TrieNode) usize {
+    var count: usize = 0;
+    if (node.eof) {
+        count += 1;
+    }
 
-    while (len < min_len) : (len += 1) {
-        if (a[len] != b[len]) {
-            break;
+    for (0..MAX_CHILDREN) |i| {
+        const child = node.children[i];
+        if (child != null) {
+            count += countKeys(child.?);
         }
     }
+    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;
+    }
+
+    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);
+        }
+    }
+}
+
+pub fn getKeysFromNode(node: *TrieNode, prefix: []const u8) [][]const u8 {
+    const keyCount = countKeys(node);
+    if (keyCount == 0) {
+        return &[_][]const u8{};
+    }
+
+    const keys = allocator.alloc([]const u8, keyCount) catch unreachable;
+    var index: usize = 0;
+    collectKeys(node, prefix, keys, &index);
+    return keys;
+}
+
+pub fn getKeysByPrefix(prefix: []const u8) []const u8 {
+    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;
+}
+
+pub fn getValuesByPrefix(prefix: []const u8) []const u8 {
+    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;
+    for (keys, 0..) |key, i| {
+        const value = storage.read(key) orelse "";
+        values[i] = value;
+    }
+
+    return std.mem.join(allocator, "\r", values) catch unreachable;
+}
 
-    return len;
+pub fn getAllKeys() []const u8 {
+    const keys = getKeysFromNode(&root, &[_]u8{});
+    if (keys.len == 0) return "";
+    return std.mem.join(allocator, "\r", keys) catch unreachable;
 }

main.zig

@@ -15,7 +15,7 @@ pub fn main() !void {
     defer posix.close(listener);
 
     std.debug.print("2025 pizzakv! TCP Listening on port {any}\n<danilo@fragoso.dev>\n---------\n", .{PORT});
-    std.debug.print("Commands:\n\nread key\nwrite key|value\ndelete key\nstatus\n", .{});
+    std.debug.print("Commands:\n\nread key\nwrite key|value\ndelete key\nkeys\nreads prefix\nstatus\n", .{});
     std.debug.print("---------\n", .{});
 
     try persistence.init();
@@ -38,7 +38,7 @@ pub fn handleConnection(conn: posix.socket_t) !void {
     defer posix.close(conn);
 
     var requestBuffer: [1024 * 1024]u8 = undefined;
-    var responseBuffer: [1024 * 1024]u8 = undefined;
+    var responseBuffer: [1024 * 1024 * 4]u8 = undefined;
 
     while (true) {
         const n = try socket.readUntilCR(conn, &requestBuffer);

storage.zig

@@ -1,4 +1,6 @@
 const std = @import("std");
+
+const index = @import("index.zig");
 const hashing = @import("hashing.zig");
 const persistence = @import("persistence.zig");
 
@@ -18,9 +20,9 @@ var mutex: std.Thread.Mutex = .{};
 
 pub fn writeVolatile(key: []const u8, value: []const u8) ?*Entry {
     const hash = hashing.hashKey(key);
-    const index = hash % buf.len;
+    const bufIdx = hash % buf.len;
 
-    var current = buf[index];
+    var current = buf[bufIdx];
     while (current) |entry| {
         if (std.mem.eql(u8, entry.key, key)) {
             allocator.free(entry.value);
@@ -36,10 +38,11 @@ pub fn writeVolatile(key: []const u8, value: []const u8) ?*Entry {
     newEntry.* = Entry{
         .key = allocator.dupe(u8, key) catch return null,
         .value = allocator.dupe(u8, value) catch return null,
-        .next = buf[index],
+        .next = buf[bufIdx],
     };
 
-    buf[index] = newEntry;
+    buf[bufIdx] = newEntry;
+    index.insert(key);
     return newEntry;
 }
 
@@ -74,9 +77,9 @@ pub fn read(key: []const u8) ?[]const u8 {
 
 pub fn deleteVolatile(key: []const u8) bool {
     const hash = hashing.hashKey(key);
-    const index = hash % buf.len;
+    const bufIdx = hash % buf.len;
 
-    var current = buf[index];
+    var current = buf[bufIdx];
     var prev: ?*Entry = null;
 
     while (current) |entry| {
@@ -84,9 +87,10 @@ pub fn deleteVolatile(key: []const u8) bool {
             if (prev) |p| {
                 p.next = entry.next;
             } else {
-                buf[index] = entry.next;
+                buf[bufIdx] = entry.next;
             }
 
+            index.delete(key);
             allocator.free(entry.key);
             allocator.free(entry.value);
             allocator.destroy(entry);