From b7a8c8a62286eba9b612a3e000d612fe30a71ee6 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Wed, 2 Sep 2026 12:45:26 +0200 Subject: [PATCH] feat: add scoreBound and add to collections --- src/collections.zig | 156 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/src/collections.zig b/src/collections.zig index 1c0157b..c9214f6 100644 --- a/src/collections.zig +++ b/src/collections.zig @@ -791,3 +791,159 @@ pub const ZSet = struct { return out.toOwnedSlice(); } + /// Score-index key truncated to just the ordered score (no member + /// suffix), used as a scan bound. Every real entry is a strict byte + /// extension of this (the member part is never empty), so a key equal to + /// `scoreBound(s)` never itself denotes a stored member. + fn scoreBound(allocator: Allocator, name: []const u8, scoreVal: f64) ![]u8 { + var out = std.ArrayList(u8).init(allocator); + errdefer out.deinit(); + try out.append(nsTag); + try appendSegment(&out, "Z"); + try appendSegment(&out, name); + try appendSegment(&out, "S"); + try row.appendOrdered(&out, .{ .float = scoreVal }); + return out.toOwnedSlice(); + } + + /// Upserts `member`'s score, atomically re-pointing the score index if + /// the member already existed with a different score. + pub fn add(db: *Db, allocator: Allocator, name: []const u8, member: []const u8, scoreVal: f64) !void { + const mk = try memberKey(allocator, name, member); + defer allocator.free(mk); + + var attempt: usize = 0; + while (attempt < maxConflictRetries) : (attempt += 1) { + var txn = Transaction.begin(db); + defer txn.deinit(); + + const old = try txn.get(allocator, mk); + if (old) |raw| { + defer allocator.free(raw); + const oldScore = try readF64(raw); + if (oldScore == scoreVal) return; + const oldEntryKey = try scoreEntryKey(allocator, name, oldScore, member); + defer allocator.free(oldEntryKey); + try txn.delete(oldEntryKey); + } + + var scoreBuf = std.ArrayList(u8).init(allocator); + defer scoreBuf.deinit(); + try putF64(&scoreBuf, scoreVal); + try txn.put(mk, scoreBuf.items); + + const entryKey = try scoreEntryKey(allocator, name, scoreVal, member); + defer allocator.free(entryKey); + try txn.put(entryKey, &.{}); + + txn.commit() catch |err| { + if (err == error.Conflict) continue; + return err; + }; + return; + } + return error.TooManyConflicts; + } + + pub fn score(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !?f64 { + const mk = try memberKey(allocator, name, member); + defer allocator.free(mk); + const raw = try db.get(allocator, mk) orelse return null; + defer allocator.free(raw); + return try readF64(raw); + } + + pub fn incrBy(db: *Db, allocator: Allocator, name: []const u8, member: []const u8, delta: f64) !f64 { + const mk = try memberKey(allocator, name, member); + defer allocator.free(mk); + + var attempt: usize = 0; + while (attempt < maxConflictRetries) : (attempt += 1) { + var txn = Transaction.begin(db); + defer txn.deinit(); + + const oldRaw = try txn.get(allocator, mk); + const oldScore: f64 = if (oldRaw) |raw| blk: { + defer allocator.free(raw); + break :blk try readF64(raw); + } else 0.0; + const newScore = oldScore + delta; + + if (oldRaw != null) { + const oldEntryKey = try scoreEntryKey(allocator, name, oldScore, member); + defer allocator.free(oldEntryKey); + try txn.delete(oldEntryKey); + } + + var scoreBuf = std.ArrayList(u8).init(allocator); + defer scoreBuf.deinit(); + try putF64(&scoreBuf, newScore); + try txn.put(mk, scoreBuf.items); + + const entryKey = try scoreEntryKey(allocator, name, newScore, member); + defer allocator.free(entryKey); + try txn.put(entryKey, &.{}); + + txn.commit() catch |err| { + if (err == error.Conflict) continue; + return err; + }; + return newScore; + } + return error.TooManyConflicts; + } + + /// Returns true if the member was present and has been removed. + pub fn remove(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !bool { + const mk = try memberKey(allocator, name, member); + defer allocator.free(mk); + + var attempt: usize = 0; + while (attempt < maxConflictRetries) : (attempt += 1) { + var txn = Transaction.begin(db); + defer txn.deinit(); + + const raw = try txn.get(allocator, mk) orelse return false; + defer allocator.free(raw); + const existingScore = try readF64(raw); + + try txn.delete(mk); + const entryKey = try scoreEntryKey(allocator, name, existingScore, member); + defer allocator.free(entryKey); + try txn.delete(entryKey); + + txn.commit() catch |err| { + if (err == error.Conflict) continue; + return err; + }; + return true; + } + return error.TooManyConflicts; + } + + pub fn card(db: *Db, allocator: Allocator, name: []const u8) !usize { + const p = try scorePrefix(allocator, name); + defer allocator.free(p); + var it = try db.scanPrefix(allocator, p); + defer it.deinit(); + var count: usize = 0; + while (try it.next()) |_| count += 1; + return count; + } + + fn decodeScoreEntry(allocator: Allocator, name: []const u8, key: []const u8) !ZMember { + const p = try scorePrefix(allocator, name); + defer allocator.free(p); + if (key.len <= p.len or !std.mem.startsWith(u8, key, p)) return error.Malformed; + const rest = key[p.len..]; + + // rest = [tag u8][8-byte ordered score][tag u8][escaped member][0x00 0x00] + if (rest.len < 9 or rest[0] != @intFromEnum(row.ValueTag.float)) return error.Malformed; + var bits = std.mem.readInt(u64, rest[1..9], .big); + if (bits & 0x8000000000000000 != 0) { + bits &= 0x7fffffffffffffff; + } else { + bits = ~bits; + } + const decodedScore: f64 = @bitCast(bits); +