Skip to main content

kevy_store/
hash.rs

1//! `Store` hash commands.
2
3use crate::util::parse_i64;
4use crate::value::{HashData, SmallBytes, Value, hash_field_weight};
5use crate::{Entry, Store, StoreError, now_ns};
6use std::sync::Arc;
7
8impl Store {
9    // ---- hashes --------------------------------------------------------
10
11    /// Borrow the key's hash mutably, optionally creating it. `Ok(None)` means
12    /// the key is absent and `create` was false.
13    fn hash_mut(&mut self, key: &[u8], create: bool) -> Result<Option<&mut HashData>, StoreError> {
14        if self.live_entry_mut(key).is_none() {
15            if !create {
16                return Ok(None);
17            }
18            self.insert_entry(
19                SmallBytes::from_slice(key),
20                Entry::new(Value::Hash(Arc::default()), None),
21            );
22        }
23        match &mut self.map.get_mut(key).expect("present").value {
24            Value::Hash(h) => Ok(Some(Arc::make_mut(h))),
25            _ => Err(StoreError::WrongType),
26        }
27    }
28
29    /// Read the key's hash immutably (lazily expiring). `Ok(None)` if absent.
30    fn hash_ref(&mut self, key: &[u8]) -> Result<Option<&HashData>, StoreError> {
31        match self.live_entry(key) {
32            None => Ok(None),
33            Some(e) => match &e.value {
34                Value::Hash(h) => Ok(Some(h.as_ref())),
35                _ => Err(StoreError::WrongType),
36            },
37        }
38    }
39
40    /// `HSET` — returns the count of newly-added fields.
41    pub fn hset(&mut self, key: &[u8], pairs: &[(Vec<u8>, Vec<u8>)]) -> Result<usize, StoreError> {
42        let (added, delta) = {
43            let h = self.hash_mut(key, true)?.expect("created");
44            let mut a = 0usize;
45            let mut d: i64 = 0;
46            for (f, v) in pairs {
47                let smb = SmallBytes::from_slice(f);
48                let new_w = hash_field_weight(&smb, v.len()) as i64;
49                match h.insert(smb, v.clone()) {
50                    None => {
51                        a += 1;
52                        d += new_w;
53                    }
54                    Some(old) => {
55                        d += v.len() as i64 - old.len() as i64;
56                    }
57                }
58            }
59            (a, d)
60        };
61        self.account_delta(key, delta);
62        Ok(added)
63    }
64
65    /// `HSETNX` — set only if the field is absent; returns whether it was set.
66    pub fn hsetnx(&mut self, key: &[u8], field: &[u8], val: &[u8]) -> Result<bool, StoreError> {
67        let outcome = {
68            let h = self.hash_mut(key, true)?.expect("created");
69            if h.contains_key(field) {
70                if h.is_empty() {
71                    HsetnxOutcome::DropEmpty
72                } else {
73                    HsetnxOutcome::AlreadyExists
74                }
75            } else {
76                let smb = SmallBytes::from_slice(field);
77                let w = hash_field_weight(&smb, val.len()) as i64;
78                h.insert(smb, val.to_vec());
79                HsetnxOutcome::Inserted(w)
80            }
81        };
82        match outcome {
83            HsetnxOutcome::DropEmpty => {
84                self.remove_entry(key);
85                Ok(false)
86            }
87            HsetnxOutcome::AlreadyExists => Ok(false),
88            HsetnxOutcome::Inserted(w) => {
89                self.account_delta(key, w);
90                Ok(true)
91            }
92        }
93    }
94
95    pub fn hget(&mut self, key: &[u8], field: &[u8]) -> Result<Option<&[u8]>, StoreError> {
96        Ok(self
97            .hash_ref(key)?
98            .and_then(|h| h.get(field))
99            .map(std::vec::Vec::as_slice))
100    }
101
102    pub fn hexists(&mut self, key: &[u8], field: &[u8]) -> Result<bool, StoreError> {
103        Ok(self.hash_ref(key)?.is_some_and(|h| h.contains_key(field)))
104    }
105
106    pub fn hlen(&mut self, key: &[u8]) -> Result<usize, StoreError> {
107        Ok(self.hash_ref(key)?.map_or(0, kevy_map::KevyMap::len))
108    }
109
110    pub fn hmget(
111        &mut self,
112        key: &[u8],
113        fields: &[Vec<u8>],
114    ) -> Result<Vec<Option<Vec<u8>>>, StoreError> {
115        let h = self.hash_ref(key)?;
116        Ok(fields
117            .iter()
118            .map(|f| h.and_then(|h| h.get(f.as_slice())).cloned())
119            .collect())
120    }
121
122    /// `HGETALL` — flat `[field, value, field, value, ...]` (clones; perf-polish later).
123    pub fn hgetall(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
124        match self.hash_ref(key)? {
125            None => Ok(Vec::new()),
126            Some(h) => {
127                let mut out = Vec::with_capacity(h.len() * 2);
128                for (f, v) in h {
129                    out.push(f.to_vec());
130                    out.push(v.clone());
131                }
132                Ok(out)
133            }
134        }
135    }
136
137    pub fn hkeys(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
138        Ok(self
139            .hash_ref(key)?
140            .map_or(Vec::new(), |h| h.keys().map(kevy_bytes::SmallBytes::to_vec).collect()))
141    }
142
143    pub fn hvals(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
144        Ok(self
145            .hash_ref(key)?
146            .map_or(Vec::new(), |h| h.values().cloned().collect()))
147    }
148
149    /// `HDEL` — returns count removed; deletes the key if the hash becomes empty.
150    pub fn hdel(&mut self, key: &[u8], fields: &[Vec<u8>]) -> Result<usize, StoreError> {
151        let now = now_ns();
152        if !self.reap(key, now) {
153            return Ok(0);
154        }
155        let (removed, delta, drop_key) = {
156            let h_entry = self.map.get_mut(key).expect("live");
157            match &mut h_entry.value {
158                Value::Hash(h) => {
159                    let h = Arc::make_mut(h);
160                    let mut r = 0usize;
161                    let mut d: i64 = 0;
162                    for f in fields {
163                        if let Some(old_v) = h.remove(f.as_slice()) {
164                            r += 1;
165                            // The field key matters as a SmallBytes only for
166                            // heap_bytes/slot overhead; reconstruct the same
167                            // weight figure that hset paid in.
168                            let smb = SmallBytes::from_slice(f);
169                            d -= hash_field_weight(&smb, old_v.len()) as i64;
170                        }
171                    }
172                    let drop_now = h.is_empty();
173                    (r, d, drop_now)
174                }
175                _ => return Err(StoreError::WrongType),
176            }
177        };
178        if drop_key {
179            self.remove_entry(key);
180        } else {
181            self.account_delta(key, delta);
182        }
183        Ok(removed)
184    }
185
186    /// `HINCRBY` — preserves TTL; errors if the field isn't an integer.
187    pub fn hincrby(&mut self, key: &[u8], field: &[u8], delta: i64) -> Result<i64, StoreError> {
188        let (next, weight_delta) = {
189            let h = self.hash_mut(key, true)?.expect("created");
190            let cur = match h.get(field) {
191                Some(v) => parse_i64(v).ok_or(StoreError::NotInteger)?,
192                None => 0,
193            };
194            let next = cur.checked_add(delta).ok_or(StoreError::Overflow)?;
195            let new_bytes = next.to_string().into_bytes();
196            let smb = SmallBytes::from_slice(field);
197            let new_field_w = hash_field_weight(&smb, new_bytes.len()) as i64;
198            let new_value_len = new_bytes.len();
199            let wd = match h.insert(smb, new_bytes) {
200                None => new_field_w,
201                Some(old) => new_value_len as i64 - old.len() as i64,
202            };
203            (next, wd)
204        };
205        self.account_delta(key, weight_delta);
206        Ok(next)
207    }
208}
209
210enum HsetnxOutcome {
211    DropEmpty,
212    AlreadyExists,
213    Inserted(i64),
214}