From 025f1708c6956e67677340428000f70ba3a95c5a Mon Sep 17 00:00:00 2001 From: DualFroz Date: Wed, 2 Sep 2026 14:31:02 +0200 Subject: [PATCH] feat: add rangeByRank and rangeByScore to collections --- src/collections.zig | 158 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 158 insertions(+) diff --git a/src/collections.zig b/src/collections.zig index c9214f6..6a1abb5 100644 --- a/src/collections.zig +++ b/src/collections.zig @@ -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()); +} +