feat: extend sstable with commonPrefixLen and seekInBlock

src/sstable.zig

@@ -24,11 +24,26 @@ fn order(a: []const u8, b: []const u8) std.math.Order {
 }
 
 const magic: u32 = 0xC0BA1701;
-const version: u32 = 1;
+const version: u32 = 2;
 const footerLen = 4 + 4 + 8 + 4 + 8 + 4 + 8 + 8;
 const recordHeader = 1 + 8 + 4 + 4;
 const defaultBlockTarget = 4096;
 
+/// Every Nth record in a data block is a "restart point": stored with its
+/// full key (shared == 0) instead of a shared-prefix delta against the
+/// previous key. This bounds how many deltas must be re-applied to recover
+/// any one key (never more than `restartInterval - 1`) and, more
+/// importantly here, gives block expansion a cheap, structural corruption
+/// check: a real restart offset must always land on a shared == 0 record.
+const restartInterval: usize = 16;
+
+fn commonPrefixLen(a: []const u8, b: []const u8) usize {
+    const n = @min(a.len, b.len);
+    var i: usize = 0;
+    while (i < n and a[i] == b[i]) : (i += 1) {}
+    return i;
+}
+
 pub const Builder = struct {
     allocator: Allocator,
     file: std.fs.File,
@@ -36,6 +51,9 @@ pub const Builder = struct {
     blockBuf: std.ArrayList(u8),
     blockFirstKey: std.ArrayList(u8),
     indexBuf: std.ArrayList(u8),
+    lastKey: std.ArrayList(u8),
+    blockRestarts: std.ArrayList(u32),
+    blockRecordCount: usize,
     numEntries: u64,
     logicalBytes: u64,
     maxSeq: u64,
@@ -51,6 +69,9 @@ pub const Builder = struct {
             .blockBuf = std.ArrayList(u8).init(allocator),
             .blockFirstKey = std.ArrayList(u8).init(allocator),
             .indexBuf = std.ArrayList(u8).init(allocator),
+            .lastKey = std.ArrayList(u8).init(allocator),
+            .blockRestarts = std.ArrayList(u32).init(allocator),
+            .blockRecordCount = 0,
             .numEntries = 0,
             .logicalBytes = 0,
             .maxSeq = 0,
@@ -63,6 +84,8 @@ pub const Builder = struct {
         self.blockBuf.deinit();
         self.blockFirstKey.deinit();
         self.indexBuf.deinit();
+        self.lastKey.deinit();
+        self.blockRestarts.deinit();
         self.bloom.deinit();
         self.file.close();
     }
@@ -72,14 +95,25 @@ pub const Builder = struct {
             self.blockFirstKey.clearRetainingCapacity();
             try self.blockFirstKey.appendSlice(record.key);
         }
+
+        const isRestart = self.blockRecordCount % restartInterval == 0;
+        const shared: u32 = if (isRestart) 0 else @intCast(commonPrefixLen(self.lastKey.items, record.key));
+        const unshared = record.key[shared..];
+        if (isRestart) try self.blockRestarts.append(@intCast(self.blockBuf.items.len));
+
         const writer = encode.ByteWriter.init(&self.blockBuf);
         try writer.byte(@intFromEnum(record.kind));
         try writer.putU64(record.seq);
-        try writer.putU32(@intCast(record.key.len));
+        try writer.putU32(shared);
+        try writer.putU32(@intCast(unshared.len));
         try writer.putU32(@intCast(record.value.len));
-        try writer.slice(record.key);
+        try writer.slice(unshared);
         try writer.slice(record.value);
 
+        self.lastKey.clearRetainingCapacity();
+        try self.lastKey.appendSlice(record.key);
+        self.blockRecordCount += 1;
+
         self.bloom.add(record.key);
         self.numEntries += 1;
         self.logicalBytes += recordHeader + record.key.len + record.value.len;
@@ -90,6 +124,15 @@ pub const Builder = struct {
 
     fn flushBlock(self: *Builder) !void {
         if (self.blockBuf.items.len == 0) return;
+
+        // The restart-offset trailer rides inside the same logical block
+        // bytes as the records themselves, so it survives compression and
+        // is immediately available to `expandBlock` on the read side
+        // without a separate on-disk section or footer entry.
+        const trailerWriter = encode.ByteWriter.init(&self.blockBuf);
+        for (self.blockRestarts.items) |restartOffset| try trailerWriter.putU32(restartOffset);
+        try trailerWriter.putU32(@intCast(self.blockRestarts.items.len));
+
         const originalLen: u32 = @intCast(self.blockBuf.items.len);
 
         const encoded = try compress.encode(self.allocator, self.blockBuf.items);
@@ -116,6 +159,9 @@ pub const Builder = struct {
 
         self.offset += physLen + 4;
         self.blockBuf.clearRetainingCapacity();
+        self.blockRestarts.clearRetainingCapacity();
+        self.blockRecordCount = 0;
+        self.lastKey.clearRetainingCapacity();
     }
 
     pub fn finish(self: *Builder) !void {
@@ -160,6 +206,13 @@ const IndexEntry = struct {
     offset: u64,
     blockLen: u32,
     firstKey: []const u8,
+    // Offsets (relative to this block's own bytes, i.e. usable directly
+    // against the slice `blockBytes()` returns) of every restart record in
+    // the flattened block, in ascending order. restarts[0] is always 0.
+    // Lets point lookups binary-search straight to the <= restartInterval
+    // records that could hold the key instead of linear-scanning the whole
+    // block from position 0.
+    restarts: []u32,
 };
 
 const blockPrefix = 1 + 4;
@@ -188,6 +241,7 @@ pub const Reader = struct {
         if (buffer.len < footerLen) return error.CorruptSstable;
         const footerStart = buffer.len - footerLen;
         if (readU32(buffer, footerStart) != magic) return error.CorruptSstable;
+        if (readU32(buffer, footerStart + 4) != version) return error.UnsupportedSstableVersion;
         const indexOffset = readU64(buffer, footerStart + 8);
         const indexLen = readU32(buffer, footerStart + 16);
         const bloomOffset = readU64(buffer, footerStart + 20);
@@ -204,7 +258,10 @@ pub const Reader = struct {
         var decoded = std.ArrayList(u8).init(allocator);
         errdefer decoded.deinit();
         var entries = std.ArrayList(IndexEntry).init(allocator);
-        errdefer entries.deinit();
+        errdefer {
+            for (entries.items) |e| allocator.free(e.restarts);
+            entries.deinit();
+        }
 
         var pos: usize = 0;
         while (pos < indexData.len) {
@@ -221,11 +278,19 @@ pub const Reader = struct {
             const payload = physical[blockPrefix..];
             const plainBlock = try compress.decode(allocator, codec, payload, originalLen);
             defer allocator.free(plainBlock);
+            const expanded = try expandBlock(allocator, plainBlock);
+            defer allocator.free(expanded.bytes);
+            errdefer allocator.free(expanded.restarts);
 
             const decOffset = decoded.items.len;
-            try decoded.appendSlice(plainBlock);
+            try decoded.appendSlice(expanded.bytes);
 
-            try entries.append(.{ .offset = decOffset, .blockLen = @intCast(plainBlock.len), .firstKey = firstKey });
+            try entries.append(.{
+                .offset = decOffset,
+                .blockLen = @intCast(expanded.bytes.len),
+                .firstKey = firstKey,
+                .restarts = expanded.restarts,
+            });
             pos = keyStart + firstKeyLen;
         }
 
@@ -265,6 +330,7 @@ pub const Reader = struct {
     }
 
     pub fn deinit(self: *Reader) void {
+        for (self.index) |e| self.allocator.free(e.restarts);
         self.allocator.free(self.index);
         self.bloom.deinit();
         self.allocator.free(self.dataBuf);
@@ -279,7 +345,7 @@ pub const Reader = struct {
         if (!self.bloom.mayContain(key)) return null;
         const blockIdx = self.findBlock(key) orelse return null;
         const block = self.blockBytes(blockIdx);
-        var pos: usize = 0;
+        var pos: usize = Reader.seekInBlock(block, self.index[blockIdx].restarts, key);
         while (pos < block.len) {
             const decoded = decodeRecord(block, pos);
             switch (order(decoded.record.key, key)) {
@@ -307,7 +373,7 @@ pub const Reader = struct {
         const block = cached orelse self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
         if (cached == null) try blockCache.put(blockKey, block);
 
-        var pos: usize = 0;
+        var pos: usize = Reader.seekInBlock(block, entry.restarts, key);
         while (pos < block.len) {
             const decoded = decodeRecord(block, pos);
             switch (order(decoded.record.key, key)) {
@@ -329,7 +395,10 @@ pub const Reader = struct {
     pub fn cursor(self: *const Reader, start: ?[]const u8) Cursor {
         var c = Cursor{ .reader = self, .blockIdx = 0, .pos = 0, .nextPos = 0, .cur = null };
         if (self.index.len == 0) return c;
-        c.blockIdx = if (start) |s| (self.findBlock(s) orelse 0) else 0;
+        if (start) |s| {
+            c.blockIdx = self.findBlock(s) orelse 0;
+            c.pos = Reader.seekInBlock(self.blockBytes(c.blockIdx), self.index[c.blockIdx].restarts, s);
+        }
         c.refresh();
         if (start) |s| {
             while (c.valid() and order(c.cur.?.key, s) == .lt) c.advance();
@@ -400,6 +469,30 @@ pub const Reader = struct {
         return ans;
     }
 
+    /// Binary-searches `restarts` (each entry the flat-block offset of a
+    /// restart record, ascending by key since keys are sorted within a
+    /// block) for the rightmost restart whose key is <= `key`, and returns
+    /// its offset. Callers then linear-scan forward from there instead of
+    /// from position 0, bounding the scan to at most `restartInterval`
+    /// records regardless of how many records the block holds. Returns 0
+    /// (restarts[0], always the block's first record) when `key` sorts
+    /// before every restart.
+    fn seekInBlock(block: []const u8, restarts: []const u32, key: []const u8) usize {
+        if (restarts.len == 0) return 0;
+        var lo: usize = 0;
+        var hi: usize = restarts.len;
+        while (lo < hi) {
+            const mid = lo + (hi - lo) / 2;
+            const rec = decodeRecord(block, restarts[mid]);
+            if (order(rec.record.key, key) != .gt) {
+                lo = mid + 1;
+            } else {
+                hi = mid;
+            }
+        }
+        return if (lo == 0) 0 else restarts[lo - 1];
+    }
+
     fn blockBytes(self: *const Reader, idx: usize) []const u8 {
         const entry = self.index[idx];
         return self.dataBuf[entry.offset .. entry.offset + entry.blockLen];