feat: add rangeByRank and rangeByScore to collections

This commit is contained in:
dualfroz committed 2026-09-02 14:31:02 +02:00
1 parent b7a8c8a622
commit 025f1708c6
1 file changed
+158
+158
View File
@@ -947,3 +947,161 @@ pub const ZSet = struct {
}
const decodedScore: f64 = @bitCast(bits);
const memberPart = rest[9..];
if (memberPart.len < 3 or memberPart[0] != @intFromEnum(row.ValueTag.text)) return error.Malformed;
var unescaped = std.ArrayList(u8).init(allocator);
errdefer unescaped.deinit();
var i: usize = 1;
while (i < memberPart.len) {
if (memberPart[i] == 0x00) {
if (i + 1 >= memberPart.len) return error.Malformed;
if (memberPart[i + 1] == 0x00) break;
if (memberPart[i + 1] == 0xFF) {
try unescaped.append(0x00);
i += 2;
continue;
}
return error.Malformed;
}
try unescaped.append(memberPart[i]);
i += 1;
}
return .{ .member = try unescaped.toOwnedSlice(), .score = decodedScore };
}
/// Ascending-by-score range selected by rank (0-based, inclusive on both
/// ends), Redis-style: negative indices count from the end. Caller owns
/// the returned slice and each member's buffer.
pub fn rangeByRank(db: *Db, allocator: Allocator, name: []const u8, start: i64, stop: i64) ![]ZMember {
const total: i64 = @intCast(try card(db, allocator, name));
var results = std.ArrayList(ZMember).init(allocator);
errdefer {
for (results.items) |m| m.free(allocator);
results.deinit();
}
if (total == 0) return results.toOwnedSlice();
var lo = if (start < 0) total + start else start;
var hi = if (stop < 0) total + stop else stop;
if (lo < 0) lo = 0;
if (hi > total - 1) hi = total - 1;
if (lo > hi) return results.toOwnedSlice();
const p = try scorePrefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var idx: i64 = 0;
while (try it.next()) |entry| {
if (idx >= lo and idx <= hi) {
try results.append(try decodeScoreEntry(allocator, name, entry.key));
}
idx += 1;
}
return results.toOwnedSlice();
}
/// All members whose score falls within [min, max], found via a real
/// bounded scan over the score index (not a full-collection scan): the
/// start bound is `scoreBound(min)` and the end bound is the byte
/// immediately after `scoreBound(max)`, which is enough to include every
/// entry whose score equals `max` since real entries strictly extend
/// their score bound (see `scoreBound`'s doc comment). Caller owns the
/// returned slice and each member's buffer.
pub fn rangeByScore(db: *Db, allocator: Allocator, name: []const u8, min: f64, max: f64) ![]ZMember {
var results = std.ArrayList(ZMember).init(allocator);
errdefer {
for (results.items) |m| m.free(allocator);
results.deinit();
}
const startBound = try scoreBound(allocator, name, min);
defer allocator.free(startBound);
const maxBound = try scoreBound(allocator, name, max);
defer allocator.free(maxBound);
const endBound = try prefixUpperBound(allocator, maxBound) orelse try allocator.dupe(u8, maxBound);
defer allocator.free(endBound);
var it = try db.scanRange(allocator, .{
.start = startBound,
.end = endBound,
.startInclusive = true,
.endInclusive = false,
});
defer it.deinit();
while (try it.next()) |entry| {
try results.append(try decodeScoreEntry(allocator, name, entry.key));
}
return results.toOwnedSlice();
}
/// 0-based ascending-by-score rank of `member`, or null if absent.
pub fn rank(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !?usize {
const target = try score(db, allocator, name, member) orelse return null;
const p = try scorePrefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var idx: usize = 0;
while (try it.next()) |entry| {
const decoded = try decodeScoreEntry(allocator, name, entry.key);
defer decoded.free(allocator);
if (decoded.score == target and std.mem.eql(u8, decoded.member, member)) return idx;
idx += 1;
}
return null;
}
};
// ---------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------
fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !Db {
const dbPath = try tmp.dir.realpathAlloc(allocator, ".");
defer allocator.free(dbPath);
return Db.open(allocator, dbPath, .{});
}
test "counter incrBy accumulates and get defaults to zero" {
const allocator = std.testing.allocator;
var tmp = std.testing.tmpDir(.{});
defer tmp.cleanup();
var db = try openTemp(&tmp, allocator);
defer db.close();
try std.testing.expectEqual(@as(i64, 0), try Counter.get(&db, allocator, "hits"));
try std.testing.expectEqual(@as(i64, 5), try Counter.incrBy(&db, allocator, "hits", 5));
try std.testing.expectEqual(@as(i64, 3), try Counter.incrBy(&db, allocator, "hits", -2));
try std.testing.expectEqual(@as(i64, 3), try Counter.get(&db, allocator, "hits"));
try Counter.set(&db, allocator, "hits", 100);
try std.testing.expectEqual(@as(i64, 100), try Counter.get(&db, allocator, "hits"));
}
test "counter incrBy rejects a real external conflict" {
const allocator = std.testing.allocator;
var tmp = std.testing.tmpDir(.{});
defer tmp.cleanup();
var db = try openTemp(&tmp, allocator);
defer db.close();
// incrBy retries internally, so exercise the underlying primitive
// directly: begin a txn, let an external writer land in between, then
// confirm the stale commit is rejected rather than silently overwriting.
const k = try Counter.key(allocator, "race");
defer allocator.free(k);
var txn = Transaction.begin(&db);
defer txn.deinit();
_ = try txn.get(allocator, k);
try db.put(k, "external");
try txn.put(k, "mine");
try std.testing.expectError(error.Conflict, txn.commit());
}