feat: extend sstable with ExpandedBlock and expandBlock

src/sstable.zig

@@ -612,6 +612,113 @@ const Decoded = struct {
     next: usize,
 };
 
+const ExpandedBlock = struct {
+    bytes: []u8,
+    // Flat-block offsets (into `bytes`) of every restart record, ascending.
+    // Consumed by `Reader.seekInBlock` for binary-search point lookups.
+    restarts: []u32,
+};
+
+/// Reconstructs a physical block (prefix-compressed records plus its
+/// restart-offset trailer, as written by `Builder`) back into the flat
+/// every-record-carries-its-full-key layout that `decodeRecord` and every
+/// cursor in this file expect. Expansion happens once, right after
+/// decompression in `Reader.open`, so `dataBuf`/the block cache keep
+/// storing fully materialized blocks and the rest of the reader never has
+/// to know keys were delta-encoded on disk - in particular this preserves
+/// the invariant `db.zig`'s k-way merge relies on, that a cursor's current
+/// key slice stays valid until that same cursor is advanced again.
+///
+/// Also validates the restart trailer against the record stream: every
+/// declared restart offset must land exactly on a record whose `shared`
+/// field is 0, and every record with `shared` 0 at that exact interval
+/// must have a matching restart entry. This catches a truncated or
+/// shifted restart table that a whole-block CRC would not (the CRC only
+/// proves the compressed bytes were not corrupted in transit/on disk, not
+/// that this structural relationship inside them still holds).
+///
+/// Also re-derives, in the *flattened* block's own coordinates, the offset
+/// of each restart record (`ExpandedBlock.restarts`) so the read side can
+/// still binary-search to a restart-sized (<= restartInterval records)
+/// linear-scan window even though the on-disk shared/unshared deltas are
+/// gone post-expansion.
+fn expandBlock(allocator: Allocator, physical: []const u8) !ExpandedBlock {
+    if (physical.len < 4) return error.CorruptSstable;
+    const numRestarts = readU32(physical, physical.len - 4);
+    const restartTableBytes = @as(usize, numRestarts) * 4;
+    if (restartTableBytes + 4 > physical.len) return error.CorruptSstable;
+    const recordsEnd = physical.len - 4 - restartTableBytes;
+    const restartTable = physical[recordsEnd .. recordsEnd + restartTableBytes];
+
+    var out = std.ArrayList(u8).init(allocator);
+    errdefer out.deinit();
+
+    var flatRestarts = std.ArrayList(u32).init(allocator);
+    errdefer flatRestarts.deinit();
+
+    var prevKey = std.ArrayList(u8).init(allocator);
+    defer prevKey.deinit();
+
+    const perRecordFixed = 1 + 8 + 4 + 4 + 4;
+    var pos: usize = 0;
+    var restartIdx: usize = 0;
+    var recordIdx: usize = 0;
+    while (pos < recordsEnd) {
+        if (pos + perRecordFixed > recordsEnd) return error.CorruptSstable;
+        const kind = physical[pos];
+        const seq = readU64(physical, pos + 1);
+        const shared = readU32(physical, pos + 9);
+        const unshared = readU32(physical, pos + 13);
+        const valLen = readU32(physical, pos + 17);
+        const keyStart = pos + perRecordFixed;
+        if (shared > prevKey.items.len) return error.CorruptSstable;
+        if (keyStart + @as(usize, unshared) + @as(usize, valLen) > recordsEnd) return error.CorruptSstable;
+
+        // A genuine block (as `Builder` writes it) marks exactly every
+        // `restartInterval`-th record (by position within the block) as a
+        // restart, in order, with shared == 0. Checking both directions -
+        // a restart-interval position must be declared, and a declared
+        // offset must fall on a restart-interval position - catches a
+        // restart table that is present but shifted, truncated, or
+        // fabricated, not just one that is structurally impossible.
+        const shouldBeRestart = recordIdx % restartInterval == 0;
+        const isDeclaredRestart = restartIdx < numRestarts and readU32(restartTable, restartIdx * 4) == pos;
+        if (shouldBeRestart != isDeclaredRestart) return error.CorruptSstable;
+        if (shouldBeRestart) {
+            if (shared != 0) return error.CorruptSstable;
+            restartIdx += 1;
+            try flatRestarts.append(@intCast(out.items.len));
+        }
+
+        const unsharedBytes = physical[keyStart .. keyStart + unshared];
+        const valStart = keyStart + unshared;
+        const value = physical[valStart .. valStart + valLen];
+
+        const fullKey = try allocator.alloc(u8, shared + unshared);
+        defer allocator.free(fullKey);
+        @memcpy(fullKey[0..shared], prevKey.items[0..shared]);
+        @memcpy(fullKey[shared..], unsharedBytes);
+
+        const writer = encode.ByteWriter.init(&out);
+        try writer.byte(kind);
+        try writer.putU64(seq);
+        try writer.putU32(@intCast(fullKey.len));
+        try writer.putU32(valLen);
+        try writer.slice(fullKey);
+        try writer.slice(value);
+
+        prevKey.clearRetainingCapacity();
+        try prevKey.appendSlice(fullKey);
+
+        pos = valStart + valLen;
+        recordIdx += 1;
+    }
+    if (pos != recordsEnd) return error.CorruptSstable;
+    if (restartIdx != numRestarts) return error.CorruptSstable;
+
+    return .{ .bytes = try out.toOwnedSlice(), .restarts = try flatRestarts.toOwnedSlice() };
+}
+
 fn decodeRecord(block: []const u8, pos: usize) Decoded {
     const kind: Kind = @enumFromInt(block[pos]);
     const seq = readU64(block, pos + 1);
@@ -697,6 +804,73 @@ test "sstable cursor seek" {
     try std.testing.expect(!beyond.valid());
 }
 
+test "restart index lets point lookups skip straight past earlier records in a block" {
+    const allocator = std.testing.allocator;
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+
+    try buildFixture(allocator, tmp.dir, "test.sst", 500);
+    var reader = try Reader.open(allocator, tmp.dir, "test.sst", 1);
+    defer reader.deinit();
+
+    // Every block must have carried its restart offsets through expansion:
+    // the first is always the block's own first record, and multi-restart
+    // blocks (any block with more than restartInterval records) must have
+    // strictly increasing keys at each restart, matching overall key order.
+    var sawMultiRestartBlock = false;
+    var blockIdx: usize = 0;
+    while (blockIdx < reader.index.len) : (blockIdx += 1) {
+        const entry = reader.index[blockIdx];
+        try std.testing.expect(entry.restarts.len >= 1);
+        try std.testing.expectEqual(@as(u32, 0), entry.restarts[0]);
+        if (entry.restarts.len > 1) sawMultiRestartBlock = true;
+
+        const block = reader.blockBytes(blockIdx);
+        var prevKey: []const u8 = &.{};
+        for (entry.restarts) |r| {
+            const rec = decodeRecord(block, r);
+            if (prevKey.len > 0) try std.testing.expect(order(prevKey, rec.record.key) == .lt);
+            prevKey = rec.record.key;
+        }
+    }
+    // 500 records at ~87/block (blockTarget 4096 / ~47 bytes/record) and
+    // restartInterval 16 guarantees at least one block with several
+    // restarts - if this ever goes false the fixture stopped exercising
+    // the multi-restart path the rest of this test relies on.
+    try std.testing.expect(sawMultiRestartBlock);
+
+    // A probe exactly equal to a later restart's key must land on that
+    // restart's own offset (the binary search must not undershoot to an
+    // earlier restart when there is an exact match).
+    const entry0 = reader.index[0];
+    try std.testing.expect(entry0.restarts.len > 2);
+    const block0 = reader.blockBytes(0);
+    const secondRestartOffset = entry0.restarts[1];
+    const secondRestartKey = decodeRecord(block0, secondRestartOffset).record.key;
+    try std.testing.expectEqual(secondRestartOffset, Reader.seekInBlock(block0, entry0.restarts, secondRestartKey));
+
+    // A probe strictly between the first and second restart's keys must
+    // stay on the first restart (must not overshoot to the second).
+    const firstRestartKey = decodeRecord(block0, entry0.restarts[0]).record.key;
+    var midBuf: [16]u8 = undefined;
+    @memcpy(midBuf[0..firstRestartKey.len], firstRestartKey);
+    midBuf[firstRestartKey.len - 1] += 1; // one key past first restart's key
+    try std.testing.expect(order(midBuf[0..firstRestartKey.len], secondRestartKey) == .lt);
+    try std.testing.expectEqual(entry0.restarts[0], Reader.seekInBlock(block0, entry0.restarts, midBuf[0..firstRestartKey.len]));
+
+    // A probe past every restart's key must resolve to the last restart.
+    const lastRestartKey = decodeRecord(block0, entry0.restarts[entry0.restarts.len - 1]).record.key;
+    var pastBuf: [16]u8 = undefined;
+    @memcpy(pastBuf[0..lastRestartKey.len], lastRestartKey);
+    pastBuf[lastRestartKey.len - 1] += 1;
+    try std.testing.expectEqual(entry0.restarts[entry0.restarts.len - 1], Reader.seekInBlock(block0, entry0.restarts, pastBuf[0..lastRestartKey.len]));
+
+    // End to end: lookups still resolve correctly for keys sitting right at
+    // and around restart boundaries, now that the scan starts mid-block.
+    try std.testing.expectEqualStrings("payload-value", reader.find(firstRestartKey).?.value);
+    try std.testing.expectEqualStrings("payload-value", reader.find(secondRestartKey).?.value);
+}
+
 test "sstable reverse cursor descends fully" {
     const allocator = std.testing.allocator;
     var tmp = std.testing.tmpDir(.{});