codecrafters submit [skip ci]

5 files changed, 219 insertions, 119 deletions

diff --git a/src/main.zig b/src/main.zig
index cccdf44..780c70c 100755
--- a/src/main.zig
+++ b/src/main.zig
@@ -1,32 +1,13 @@
 const std = @import("std");
+const varint = @import("varint.zig");
+const record = @import("record.zig");
+const schema = @import("schema.zig");
+const tables = @import("tables.zig");
+
 var stdout_buffer: [1024]u8 = undefined;
 var stdout_writer = std.fs.File.stdout().writer(&stdout_buffer);
 const stdout = &stdout_writer.interface;
 
-fn parseVarint(data: []const u8) struct { value: u64, len: usize } {
-    var result: u64 = 0;
-    var i: usize = 0;
-
-    while (i < data.len and i < 9) : (i += 1) {
-        const byte = data[i];
-        result |= @as(u64, byte & 0x7f) << @as(u6, @intCast(i * 7));
-
-        if ((byte & 0x80) == 0) {
-            return .{ .value = result, .len = i + 1 };
-        }
-    }
-
-    return .{ .value = result, .len = i };
-}
-
-fn readStringFromRecord(data: []const u8, serial_type: u64) struct { value: []const u8, len: usize } {
-    if (serial_type >= 13 and (serial_type % 2) == 1) {
-        const size = (serial_type - 13) / 2;
-        return .{ .value = data[0..size], .len = size };
-    }
-    return .{ .value = "", .len = 0 };
-}
-
 pub fn main() !void {
     var gpa = std.heap.GeneralPurposeAllocator(.{}){};
     defer _ = gpa.deinit();
@@ -40,109 +21,34 @@ pub fn main() !void {
         std.process.exit(1);
     }
 
-    const database_file_path: []const u8 = args[1];
-    const command: []const u8 = args[2];
+    var file = try std.fs.cwd().openFile(args[1], .{});
+    defer file.close();
 
-    if (std.mem.eql(u8, command, ".dbinfo")) {
-        var file = try std.fs.cwd().openFile(database_file_path, .{});
-        defer file.close();
-
-        var buf: [2]u8 = undefined;
-        _ = try file.seekTo(16);
-        _ = try file.read(&buf);
-        const page_size = std.mem.readInt(u16, &buf, .big);
-        try stdout.print("database page size: {}\n", .{page_size});
+    var buf: [2]u8 = undefined;
+    _ = try file.seekTo(16);
+    _ = try file.read(&buf);
+    const page_size = std.mem.readInt(u16, &buf, .big);
 
+    if (std.mem.eql(u8, args[2], ".dbinfo")) {
         _ = try file.seekTo(103);
         _ = try file.read(&buf);
         const num_tables = std.mem.readInt(u16, &buf, .big);
+        try stdout.print("database page size: {}\n", .{page_size});
         try stdout.print("number of tables: {}\n", .{num_tables});
-        try stdout.flush();
-    } else if (std.mem.eql(u8, command, ".tables")) {
-        var file = try std.fs.cwd().openFile(database_file_path, .{});
-        defer file.close();
-
-        var buf: [2]u8 = undefined;
-        _ = try file.seekTo(16);
-        _ = try file.read(&buf);
-        const page_size = std.mem.readInt(u16, &buf, .big);
-
-        _ = try file.seekTo(103);
-        _ = try file.read(&buf);
-        const num_cells = std.mem.readInt(u16, &buf, .big);
-
-        var cell_pointers = try allocator.alloc(u16, num_cells);
-        defer allocator.free(cell_pointers);
-
-        for (0..num_cells) |i| {
-            _ = try file.seekTo(108 + i * 2);
-            _ = try file.read(&buf);
-            cell_pointers[i] = std.mem.readInt(u16, &buf, .big);
+    } else if (std.mem.eql(u8, args[2], ".tables")) {
+        try tables.showTables(allocator, &file, page_size, stdout);
+    } else {
+        var tokens = std.mem.tokenizeScalar(u8, args[2], ' ');
+        var last_token: []const u8 = "";
+        while (tokens.next()) |token| {
+            last_token = token;
         }
 
-        var page_data = try allocator.alloc(u8, page_size);
-        defer allocator.free(page_data);
-
-        _ = try file.seekTo(0);
-        _ = try file.read(page_data);
-
-        var first_table = true;
-
-        for (0..num_cells) |i| {
-            const cell_offset = cell_pointers[i];
-            const cell_data = page_data[cell_offset..];
-
-            var parsed = parseVarint(cell_data);
-            var pos = parsed.len;
-
-            parsed = parseVarint(cell_data[pos..]);
-            pos += parsed.len;
-
-            const record_start = pos;
-            const record_data = cell_data[record_start..];
-
-            parsed = parseVarint(record_data);
-            const header_size = parsed.value;
-            var header_pos = parsed.len;
+        const rootpage = try schema.getRootpage(allocator, &file, page_size, last_token);
+        const row_count = try schema.countRows(allocator, &file, page_size, rootpage);
 
-            var serial_types: [5]u64 = undefined;
-            for (0..5) |col| {
-                parsed = parseVarint(record_data[header_pos..]);
-                serial_types[col] = parsed.value;
-                header_pos += parsed.len;
-            }
-
-            var body_pos: usize = header_size;
-
-            for (0..2) |col| {
-                const serial_type = serial_types[col];
-                if (serial_type >= 13 and (serial_type % 2) == 1) {
-                    const size = (serial_type - 13) / 2;
-                    body_pos += size;
-                }
-            }
-
-            const tbl_name_type = serial_types[2];
-            const tbl_name = readStringFromRecord(record_data[body_pos..], tbl_name_type);
-
-            const body_pos_type: usize = header_size;
-            const type_serial = serial_types[0];
-            const type_str = readStringFromRecord(record_data[body_pos_type..], type_serial);
-
-            if (std.mem.eql(u8, type_str.value, "table")) {
-                if (!std.mem.eql(u8, tbl_name.value, "sqlite_sequence")) {
-                    if (!first_table) {
-                        try stdout.print(" ", .{});
-                    }
-                    try stdout.print("{s}", .{tbl_name.value});
-                    first_table = false;
-                }
-            }
-        }
-
-        if (!first_table) {
-            try stdout.print("\n", .{});
-        }
-        try stdout.flush();
+        try stdout.print("{}\n", .{row_count});
     }
+
+    try stdout.flush();
 }
diff --git a/src/record.zig b/src/record.zig
new file mode 100644
index 0000000..e430378
--- /dev/null
+++ b/src/record.zig
@@ -0,0 +1,27 @@
+const std = @import("std");
+
+pub fn readString(data: []const u8, serial_type: u64) struct { value: []const u8, len: usize } {
+    if (serial_type >= 13 and (serial_type % 2) == 1) {
+        const size = (serial_type - 13) / 2;
+        return .{ .value = data[0..size], .len = size };
+    }
+    return .{ .value = "", .len = 0 };
+}
+
+pub fn readInt(data: []const u8, serial_type: u64) struct { value: i64, len: usize } {
+    return switch (serial_type) {
+        1 => .{ .value = @as(i64, @as(i8, @intCast(data[0]))), .len = 1 },
+        2 => .{ .value = @as(i64, std.mem.readInt(i16, data[0..2], .big)), .len = 2 },
+        3 => blk: {
+            const b = [_]i32{ @intCast(data[0]), @intCast(data[1]), @intCast(data[2]) };
+            break :blk .{ .value = @as(i64, (b[0] << 16) | (b[1] << 8) | b[2]), .len = 3 };
+        },
+        4 => .{ .value = @as(i64, std.mem.readInt(i32, data[0..4], .big)), .len = 4 },
+        5 => blk: {
+            const b = [_]i64{ @intCast(data[0]), @intCast(data[1]), @intCast(data[2]), @intCast(data[3]), @intCast(data[4]), @intCast(data[5]) };
+            break :blk .{ .value = (b[0] << 40) | (b[1] << 32) | (b[2] << 24) | (b[3] << 16) | (b[4] << 8) | b[5], .len = 6 };
+        },
+        6 => .{ .value = std.mem.readInt(i64, data[0..8], .big), .len = 8 },
+        else => .{ .value = 0, .len = 0 },
+    };
+}
diff --git a/src/schema.zig b/src/schema.zig
new file mode 100644
index 0000000..00bf1f7
--- /dev/null
+++ b/src/schema.zig
@@ -0,0 +1,84 @@
+const std = @import("std");
+const varint = @import("varint.zig");
+const record = @import("record.zig");
+
+pub fn getRootpage(allocator: std.mem.Allocator, file: *std.fs.File, page_size: u16, table_name: []const u8) !u32 {
+    var buf: [2]u8 = undefined;
+    _ = try file.seekTo(103);
+    _ = try file.read(&buf);
+    const num_cells = std.mem.readInt(u16, &buf, .big);
+
+    var cell_pointers = try allocator.alloc(u16, num_cells);
+    defer allocator.free(cell_pointers);
+
+    for (0..num_cells) |i| {
+        _ = try file.seekTo(108 + i * 2);
+        _ = try file.read(&buf);
+        cell_pointers[i] = std.mem.readInt(u16, &buf, .big);
+    }
+
+    var page_data = try allocator.alloc(u8, page_size);
+    defer allocator.free(page_data);
+
+    _ = try file.seekTo(0);
+    _ = try file.read(page_data);
+
+    for (0..num_cells) |i| {
+        const cell_data = page_data[cell_pointers[i]..];
+
+        var parsed = varint.parse(cell_data);
+        var pos = parsed.len;
+
+        parsed = varint.parse(cell_data[pos..]);
+        pos += parsed.len;
+
+        const record_data = cell_data[pos..];
+        parsed = varint.parse(record_data);
+        const header_size = parsed.value;
+        var header_pos = parsed.len;
+
+        var serial_types: [5]u64 = undefined;
+        for (0..5) |col| {
+            parsed = varint.parse(record_data[header_pos..]);
+            serial_types[col] = parsed.value;
+            header_pos += parsed.len;
+        }
+
+        var body_pos: usize = header_size;
+
+        for (0..2) |col| {
+            const st = serial_types[col];
+            if (st >= 13 and (st % 2) == 1) {
+                body_pos += (st - 13) / 2;
+            }
+        }
+
+        const tbl_name_result = record.readString(record_data[body_pos..], serial_types[2]);
+        body_pos += tbl_name_result.len;
+
+        if (std.mem.eql(u8, tbl_name_result.value, table_name)) {
+            const rp = record.readInt(record_data[body_pos..], serial_types[3]);
+            return @as(u32, @intCast(rp.value));
+        }
+    }
+
+    return 0;
+}
+
+pub fn countRows(allocator: std.mem.Allocator, file: *std.fs.File, page_size: u16, rootpage: u32) !u64 {
+    if (rootpage == 0) return 0;
+
+    const page_offset = (rootpage - 1) * @as(u64, page_size);
+    var page_data = try allocator.alloc(u8, page_size);
+    defer allocator.free(page_data);
+
+    _ = try file.seekTo(page_offset);
+    _ = try file.read(page_data);
+
+    const page_type = page_data[0];
+    if (page_type == 0x05 or page_type == 0x0d) {
+        return std.mem.readInt(u16, page_data[3..5], .big);
+    }
+
+    return 0;
+}
diff --git a/src/tables.zig b/src/tables.zig
new file mode 100644
index 0000000..c317963
--- /dev/null
+++ b/src/tables.zig
@@ -0,0 +1,69 @@
+const std = @import("std");
+const varint = @import("varint.zig");
+const record = @import("record.zig");
+
+pub fn showTables(allocator: std.mem.Allocator, file: *std.fs.File, page_size: u16, stdout: anytype) !void {
+    var buf: [2]u8 = undefined;
+    _ = try file.seekTo(103);
+    _ = try file.read(&buf);
+    const num_cells = std.mem.readInt(u16, &buf, .big);
+
+    var cell_pointers = try allocator.alloc(u16, num_cells);
+    defer allocator.free(cell_pointers);
+
+    for (0..num_cells) |i| {
+        _ = try file.seekTo(108 + i * 2);
+        _ = try file.read(&buf);
+        cell_pointers[i] = std.mem.readInt(u16, &buf, .big);
+    }
+
+    var page_data = try allocator.alloc(u8, page_size);
+    defer allocator.free(page_data);
+
+    _ = try file.seekTo(0);
+    _ = try file.read(page_data);
+
+    var first = true;
+
+    for (0..num_cells) |i| {
+        const cell_data = page_data[cell_pointers[i]..];
+
+        var parsed = varint.parse(cell_data);
+        var pos = parsed.len;
+
+        parsed = varint.parse(cell_data[pos..]);
+        pos += parsed.len;
+
+        const record_data = cell_data[pos..];
+        parsed = varint.parse(record_data);
+        const header_size = parsed.value;
+        var header_pos = parsed.len;
+
+        var serial_types: [5]u64 = undefined;
+        for (0..5) |col| {
+            parsed = varint.parse(record_data[header_pos..]);
+            serial_types[col] = parsed.value;
+            header_pos += parsed.len;
+        }
+
+        var body_pos: usize = header_size;
+
+        for (0..2) |col| {
+            const st = serial_types[col];
+            if (st >= 13 and (st % 2) == 1) {
+                body_pos += (st - 13) / 2;
+            }
+        }
+
+        const type_result = record.readString(record_data[header_size..], serial_types[0]);
+        const tbl_name = record.readString(record_data[body_pos..], serial_types[2]);
+
+        if (std.mem.eql(u8, type_result.value, "table") and !std.mem.eql(u8, tbl_name.value, "sqlite_sequence")) {
+            if (!first) try stdout.print(" ", .{});
+            try stdout.print("{s}", .{tbl_name.value});
+            first = false;
+        }
+    }
+
+    if (!first) try stdout.print("\n", .{});
+}
diff --git a/src/varint.zig b/src/varint.zig
new file mode 100644
index 0000000..45207cf
--- /dev/null
+++ b/src/varint.zig
@@ -0,0 +1,14 @@
+const std = @import("std");
+
+pub fn parse(data: []const u8) struct { value: u64, len: usize } {
+    var result: u64 = 0;
+    var i: usize = 0;
+
+    while (i < data.len and i < 9) : (i += 1) {
+        const byte = data[i];
+        result |= @as(u64, byte & 0x7f) << @as(u6, @intCast(i * 7));
+        if ((byte & 0x80) == 0) return .{ .value = result, .len = i + 1 };
+    }
+
+    return .{ .value = result, .len = i };
+}