kevy_embedded/
ops_zset_algebra.rs

1//! zset algebra facades (Redis 6.2): `ZINTERSTORE` / `ZUNIONSTORE` /
2//! `ZDIFFSTORE` / `ZINTERCARD`, plus the set-algebra `*STORE` forms
3//! (`SINTERSTORE` / `SUNIONSTORE` / `SDIFFSTORE`) that 1.6.0 shipped
4//! only as reads.
5//!
6//! Locking: sources are read under their own shard locks sequentially,
7//! then `dst` is written under its shard lock — same non-atomic
8//! window as [`Store::copy`] (a concurrent writer between the reads
9//! and the store can be observed). For a fully atomic combination use
10//! the same ops inside [`Store::atomic_all_shards`].
11//!
12//! AOF: the **effect** is logged (`DEL dst` + plain `ZADD`/`SADD` of
13//! the result), never the combination — deterministic on replay and
14//! replica-apply regardless of source state.
15
16use std::io;
17
18use kevy_store::{ZAggregate, zdiff, zinter, zintercard, zunion};
19
20/// One source key's scored members (sets contribute score 1.0).
21type ScoredInput = Vec<(Vec<u8>, f64)>;
22
23#[cfg(not(target_arch = "wasm32"))]
24use crate::replica_glue::ensure_writable;
25use crate::store::{Store, commit_write, store_err};
26
27#[cfg(target_arch = "wasm32")]
28fn ensure_writable(_s: &Store) -> io::Result<()> {
29    Ok(())
30}
31
32impl Store {
33    fn gather_scored(&self, keys: &[&[u8]]) -> io::Result<Vec<ScoredInput>> {
34        keys.iter()
35            .map(|k| self.wshard(k).store.zset_or_set_members(k).map_err(store_err))
36            .collect()
37    }
38
39    fn store_zset_result(&self, dst: &[u8], pairs: &[(Vec<u8>, f64)]) -> io::Result<usize> {
40        ensure_writable(self)?;
41        let mut g = self.wshard(dst);
42        let n = g.store.zstore_result(dst, pairs);
43        commit_write(&mut g, &[b"DEL", dst])?;
44        if !pairs.is_empty() {
45            let score_strs: Vec<Vec<u8>> =
46                pairs.iter().map(|(_, s)| format!("{s}").into_bytes()).collect();
47            let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + pairs.len() * 2);
48            argv.push(b"ZADD");
49            argv.push(dst);
50            for (i, (m, _)) in pairs.iter().enumerate() {
51                argv.push(&score_strs[i]);
52                argv.push(m);
53            }
54            commit_write(&mut g, &argv)?;
55        }
56        Ok(n)
57    }
58
59    /// `ZINTERSTORE dst keys… [WEIGHTS …] [AGGREGATE SUM|MIN|MAX]` —
60    /// returns the stored cardinality.
61    pub fn zinterstore(
62        &self,
63        dst: &[u8],
64        keys: &[&[u8]],
65        weights: Option<&[f64]>,
66        aggregate: ZAggregate,
67    ) -> io::Result<usize> {
68        let inputs = self.gather_scored(keys)?;
69        self.store_zset_result(dst, &zinter(&inputs, weights, aggregate))
70    }
71
72    /// `ZUNIONSTORE dst keys… [WEIGHTS …] [AGGREGATE …]`.
73    pub fn zunionstore(
74        &self,
75        dst: &[u8],
76        keys: &[&[u8]],
77        weights: Option<&[f64]>,
78        aggregate: ZAggregate,
79    ) -> io::Result<usize> {
80        let inputs = self.gather_scored(keys)?;
81        self.store_zset_result(dst, &zunion(&inputs, weights, aggregate))
82    }
83
84    /// `ZDIFFSTORE dst keys…` (no weights/aggregate — Redis 6.2).
85    pub fn zdiffstore(&self, dst: &[u8], keys: &[&[u8]]) -> io::Result<usize> {
86        let inputs = self.gather_scored(keys)?;
87        self.store_zset_result(dst, &zdiff(&inputs))
88    }
89
90    /// `ZINTERCARD keys… [LIMIT n]` — `limit = 0` means unlimited.
91    pub fn zintercard(&self, keys: &[&[u8]], limit: usize) -> io::Result<usize> {
92        let inputs = self.gather_scored(keys)?;
93        Ok(zintercard(&inputs, limit))
94    }
95
96    /// `SINTERSTORE dst keys…` — set-algebra store form (members only).
97    pub fn sinterstore(&self, dst: &[u8], keys: &[&[u8]]) -> io::Result<usize> {
98        let members = self.sinter(keys)?;
99        self.store_set_result(dst, &members)
100    }
101
102    /// `SUNIONSTORE dst keys…`.
103    pub fn sunionstore(&self, dst: &[u8], keys: &[&[u8]]) -> io::Result<usize> {
104        let members = self.sunion(keys)?;
105        self.store_set_result(dst, &members)
106    }
107
108    /// `SDIFFSTORE dst keys…`.
109    pub fn sdiffstore(&self, dst: &[u8], keys: &[&[u8]]) -> io::Result<usize> {
110        let members = self.sdiff(keys)?;
111        self.store_set_result(dst, &members)
112    }
113
114    fn store_set_result(&self, dst: &[u8], members: &[Vec<u8>]) -> io::Result<usize> {
115        ensure_writable(self)?;
116        let mut g = self.wshard(dst);
117        let del = [dst.to_vec()];
118        g.store.del(&del);
119        commit_write(&mut g, &[b"DEL", dst])?;
120        if members.is_empty() {
121            return Ok(0);
122        }
123        let owned: Vec<Vec<u8>> = members.to_vec();
124        let n = g.store.sadd(dst, &owned).map_err(store_err)?;
125        let mut argv: Vec<&[u8]> = Vec::with_capacity(2 + members.len());
126        argv.push(b"SADD");
127        argv.push(dst);
128        argv.extend(members.iter().map(Vec::as_slice));
129        commit_write(&mut g, &argv)?;
130        Ok(n)
131    }
132}
kevy_embedded/ops_zset_algebra.rs

kevy_embedded/
ops_zset_algebra.rs
