Skip to main content

kevy_store/
list.rs

1//! `Store` list commands.
2
3use crate::small_list::{self, PushResult, SmallListData};
4use crate::util::{norm_index, range_bounds};
5use crate::value::{ListData, SmallBytes, Value, list_item_weight};
6use crate::{Entry, Store, StoreError};
7use std::sync::Arc;
8
9impl Store {
10    // ---- lists ---------------------------------------------------------
11
12    /// Borrow the key's list mutably; promote inline → heap if needed.
13    /// `create == true` materialises a fresh empty heap list when the
14    /// key is missing (the `lset/lpop/rpop/lrem/ltrim` legacy paths).
15    fn list_mut(&mut self, key: &[u8], create: bool) -> Result<Option<&mut ListData>, StoreError> {
16        if self.live_entry_mut(key).is_none() {
17            if !create {
18                return Ok(None);
19            }
20            self.insert_entry(
21                SmallBytes::from_slice(key),
22                Entry::new(Value::List(Arc::default()), None),
23            );
24        }
25        // A.8: see hash.rs::hash_mut — promote out-of-scope, then
26        // re-borrow as the heap variant.
27        let is_inline = matches!(
28            self.map.get(key).map(|e| &e.value),
29            Some(Value::SmallListInline(_))
30        );
31        if is_inline {
32            let promoted = {
33                let e = self.map.get(key).expect("present");
34                if let Value::SmallListInline(s) = &e.value {
35                    small_list::promote(s)
36                } else {
37                    unreachable!()
38                }
39            };
40            self.map.get_mut(key).expect("present").value = Value::List(Arc::new(promoted));
41            self.reweigh_entry(key);
42        }
43        match &mut self.map.get_mut(key).expect("present").value {
44            Value::List(l) => Ok(Some(Arc::make_mut(l))),
45            _ => Err(StoreError::WrongType),
46        }
47    }
48
49    /// A.8: read the key's list slot for LPUSH/RPUSH. `WrongType` on
50    /// non-list. Returns `None` when key is absent — caller creates.
51    fn list_value_for_push(&mut self, key: &[u8]) -> Result<Option<&mut Value>, StoreError> {
52        match self.live_entry_mut(key) {
53            None => Ok(None),
54            Some(e) => match &e.value {
55                Value::List(_) | Value::SmallListInline(_) => Ok(Some(&mut e.value)),
56                _ => Err(StoreError::WrongType),
57            },
58        }
59    }
60
61    /// Remove `key` if it now holds an empty list (either encoding).
62    fn drop_if_empty_list(&mut self, key: &[u8]) {
63        let empty = match self.map.get(key).map(|e| &e.value) {
64            Some(Value::List(l)) => l.is_empty(),
65            Some(Value::SmallListInline(l)) => l.is_empty(),
66            _ => false,
67        };
68        if empty {
69            self.remove_entry(key);
70        }
71    }
72
73    /// Return the list's length (either encoding). Used by the public
74    /// push functions to compute "new length" after spending entries.
75    fn list_len(&self, key: &[u8]) -> usize {
76        match self.map.get(key).map(|e| &e.value) {
77            Some(Value::List(l)) => l.len(),
78            Some(Value::SmallListInline(l)) => l.len(),
79            _ => 0,
80        }
81    }
82
83    /// `LPUSH` — prepend each value in turn; returns the new length.
84    pub fn lpush(&mut self, key: &[u8], values: &[Vec<u8>]) -> Result<usize, StoreError> {
85        let borrowed: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect();
86        self.lpush_borrowed(key, &borrowed)
87    }
88
89    /// `RPUSH` — append each value; returns the new length.
90    pub fn rpush(&mut self, key: &[u8], values: &[Vec<u8>]) -> Result<usize, StoreError> {
91        let borrowed: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect();
92        self.rpush_borrowed(key, &borrowed)
93    }
94
95    /// G4 (v1.25): borrowed-slice `LPUSH`. A.8: encoding-switch.
96    pub fn lpush_borrowed(
97        &mut self,
98        key: &[u8],
99        values: &[&[u8]],
100    ) -> Result<usize, StoreError> {
101        if values.is_empty() {
102            return Ok(self.list_len(key));
103        }
104        let mut delta: i64 = 0;
105        for v in values {
106            delta += self.list_push_one(key, v, /* front= */ true)?;
107        }
108        self.account_delta(key, delta);
109        Ok(self.list_len(key))
110    }
111
112    /// G4 (v1.25): borrowed-slice `RPUSH`. A.8: encoding-switch.
113    pub fn rpush_borrowed(
114        &mut self,
115        key: &[u8],
116        values: &[&[u8]],
117    ) -> Result<usize, StoreError> {
118        if values.is_empty() {
119            return Ok(self.list_len(key));
120        }
121        let mut delta: i64 = 0;
122        for v in values {
123            delta += self.list_push_one(key, v, /* front= */ false)?;
124        }
125        self.account_delta(key, delta);
126        Ok(self.list_len(key))
127    }
128
129    /// Push one element, applying the encoding-switch. Returns the
130    /// per-element weight delta (zero for inline, list_item_weight for
131    /// heap). `front=true` for LPUSH, `false` for RPUSH.
132    fn list_push_one(&mut self, key: &[u8], v: &[u8], front: bool) -> Result<i64, StoreError> {
133        if self.list_value_for_push(key)?.is_none() {
134            return Ok(self.list_push_create(key, v));
135        }
136        let slot = self.list_value_for_push(key)?.expect("present and a list");
137        match slot {
138            Value::SmallListInline(s) => {
139                let push = if front { s.try_push_front(v) } else { s.try_push_back(v) };
140                match push {
141                    PushResult::Pushed => Ok(0),
142                    PushResult::NoRoom => {
143                        let mut promoted = small_list::promote(s);
144                        if front {
145                            promoted.push_front(v.to_vec());
146                        } else {
147                            promoted.push_back(v.to_vec());
148                        }
149                        *slot = Value::List(Arc::new(promoted));
150                        self.reweigh_entry(key);
151                        // Reweighed from scratch — caller's delta should
152                        // be 0 for THIS pair (already in the new weight).
153                        Ok(0)
154                    }
155                }
156            }
157            Value::List(l) => {
158                let l = Arc::make_mut(l);
159                if front {
160                    l.push_front(v.to_vec());
161                } else {
162                    l.push_back(v.to_vec());
163                }
164                Ok(list_item_weight(v.len()) as i64)
165            }
166            _ => Err(StoreError::WrongType),
167        }
168    }
169
170    /// Create a fresh entry holding one element. Inline if it fits,
171    /// else heap.
172    fn list_push_create(&mut self, key: &[u8], v: &[u8]) -> i64 {
173        if let Some(inline) = SmallListData::with_one(v) {
174            self.insert_entry(
175                SmallBytes::from_slice(key),
176                Entry::new(Value::SmallListInline(inline), None),
177            );
178            0
179        } else {
180            let mut d = std::collections::VecDeque::with_capacity(1);
181            d.push_back(v.to_vec());
182            self.insert_entry(
183                SmallBytes::from_slice(key),
184                Entry::new(Value::List(Arc::new(d)), None),
185            );
186            0
187        }
188    }
189
190    /// `LPOP` — pop up to `count` from the head (deleting emptied key).
191    pub fn lpop(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
192        // Inline → promote first if there is anything to pop; simpler
193        // than maintaining a second pop path on the packed buffer.
194        if matches!(self.map.get(key).map(|e| &e.value), Some(Value::SmallListInline(_))) {
195            self.promote_list_inline_to_heap(key);
196        }
197        let (out, delta) = {
198            let mut o = Vec::new();
199            let mut d: i64 = 0;
200            if let Some(l) = self.list_mut(key, false)? {
201                for _ in 0..count {
202                    match l.pop_front() {
203                        Some(v) => {
204                            d -= list_item_weight(v.len()) as i64;
205                            o.push(v);
206                        }
207                        None => break,
208                    }
209                }
210            }
211            (o, d)
212        };
213        self.account_delta(key, delta);
214        self.drop_if_empty_list(key);
215        Ok(out)
216    }
217
218    /// `RPOP` — pop up to `count` from the tail.
219    pub fn rpop(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
220        if matches!(self.map.get(key).map(|e| &e.value), Some(Value::SmallListInline(_))) {
221            self.promote_list_inline_to_heap(key);
222        }
223        let (out, delta) = {
224            let mut o = Vec::new();
225            let mut d: i64 = 0;
226            if let Some(l) = self.list_mut(key, false)? {
227                for _ in 0..count {
228                    match l.pop_back() {
229                        Some(v) => {
230                            d -= list_item_weight(v.len()) as i64;
231                            o.push(v);
232                        }
233                        None => break,
234                    }
235                }
236            }
237            (o, d)
238        };
239        self.account_delta(key, delta);
240        self.drop_if_empty_list(key);
241        Ok(out)
242    }
243
244    /// Force-promote an inline list at `key` to its heap variant
245    /// (no-op if already heap or absent). Used by mutating paths that
246    /// only support the heap representation (pop/lrem/lset/ltrim).
247    fn promote_list_inline_to_heap(&mut self, key: &[u8]) {
248        let needs = matches!(
249            self.map.get(key).map(|e| &e.value),
250            Some(Value::SmallListInline(_))
251        );
252        if !needs {
253            return;
254        }
255        let Some(e) = self.map.get_mut(key) else { return };
256        if let Value::SmallListInline(s) = &e.value {
257            let promoted = small_list::promote(s);
258            e.value = Value::List(Arc::new(promoted));
259        }
260        self.reweigh_entry(key);
261    }
262
263    pub fn llen(&mut self, key: &[u8]) -> Result<usize, StoreError> {
264        match self.live_entry(key) {
265            None => Ok(0),
266            Some(e) => match &e.value {
267                Value::List(l) => Ok(l.len()),
268                Value::SmallListInline(l) => Ok(l.len()),
269                _ => Err(StoreError::WrongType),
270            },
271        }
272    }
273
274    pub fn lindex(&mut self, key: &[u8], idx: i64) -> Result<Option<Vec<u8>>, StoreError> {
275        match self.live_entry(key) {
276            None => Ok(None),
277            Some(e) => match &e.value {
278                Value::List(l) => Ok(norm_index(idx, l.len()).and_then(|i| l.get(i).cloned())),
279                Value::SmallListInline(l) => {
280                    let n = l.len();
281                    let Some(i) = norm_index(idx, n) else { return Ok(None) };
282                    Ok(l.iter().nth(i).map(<[u8]>::to_vec))
283                }
284                _ => Err(StoreError::WrongType),
285            },
286        }
287    }
288
289    pub fn lrange(
290        &mut self,
291        key: &[u8],
292        start: i64,
293        stop: i64,
294    ) -> Result<Vec<Vec<u8>>, StoreError> {
295        match self.live_entry(key) {
296            None => Ok(Vec::new()),
297            Some(e) => match &e.value {
298                Value::List(l) => Ok(match range_bounds(start, stop, l.len()) {
299                    None => Vec::new(),
300                    Some((s, end)) => l.iter().skip(s).take(end - s + 1).cloned().collect(),
301                }),
302                Value::SmallListInline(l) => Ok(match range_bounds(start, stop, l.len()) {
303                    None => Vec::new(),
304                    Some((s, end)) => l
305                        .iter()
306                        .skip(s)
307                        .take(end - s + 1)
308                        .map(<[u8]>::to_vec)
309                        .collect(),
310                }),
311                _ => Err(StoreError::WrongType),
312            },
313        }
314    }
315
316    /// `LSET` — errors with `NoSuchKey` / `OutOfRange` like Redis.
317    pub fn lset(&mut self, key: &[u8], idx: i64, val: &[u8]) -> Result<(), StoreError> {
318        self.promote_list_inline_to_heap(key);
319        let delta = {
320            let l = self.list_mut(key, false)?.ok_or(StoreError::NoSuchKey)?;
321            let i = norm_index(idx, l.len()).ok_or(StoreError::OutOfRange)?;
322            let old_len = l[i].len() as i64;
323            l[i] = val.to_vec();
324            val.len() as i64 - old_len
325        };
326        self.account_delta(key, delta);
327        Ok(())
328    }
329
330    /// `LREM` — remove `count` occurrences (>0 head, <0 tail, 0 all).
331    pub fn lrem(&mut self, key: &[u8], count: i64, val: &[u8]) -> Result<usize, StoreError> {
332        self.promote_list_inline_to_heap(key);
333        let (removed, delta) = {
334            let mut r = 0usize;
335            let mut d: i64 = 0;
336            match self.list_mut(key, false)? {
337                None => (0, 0),
338                Some(l) => {
339                    if count >= 0 {
340                        let limit = if count == 0 {
341                            usize::MAX
342                        } else {
343                            count as usize
344                        };
345                        let mut i = 0;
346                        while i < l.len() {
347                            if r < limit && l[i] == val {
348                                d -= list_item_weight(l[i].len()) as i64;
349                                l.remove(i);
350                                r += 1;
351                            } else {
352                                i += 1;
353                            }
354                        }
355                    } else {
356                        let limit = (-count) as usize;
357                        let mut i = l.len();
358                        while i > 0 {
359                            i -= 1;
360                            if r < limit && l[i] == val {
361                                d -= list_item_weight(l[i].len()) as i64;
362                                l.remove(i);
363                                r += 1;
364                            }
365                        }
366                    }
367                    (r, d)
368                }
369            }
370        };
371        self.account_delta(key, delta);
372        self.drop_if_empty_list(key);
373        Ok(removed)
374    }
375
376    /// `LTRIM` — keep only `[start, stop]` (deleting emptied key).
377    pub fn ltrim(&mut self, key: &[u8], start: i64, stop: i64) -> Result<(), StoreError> {
378        self.promote_list_inline_to_heap(key);
379        let delta = {
380            let mut d: i64 = 0;
381            if let Some(l) = self.list_mut(key, false)? {
382                match range_bounds(start, stop, l.len()) {
383                    None => {
384                        for v in l.iter() {
385                            d -= list_item_weight(v.len()) as i64;
386                        }
387                        l.clear();
388                    }
389                    Some((s, e)) => {
390                        for v in l.iter().skip(e + 1) {
391                            d -= list_item_weight(v.len()) as i64;
392                        }
393                        l.drain(e + 1..);
394                        for v in l.iter().take(s) {
395                            d -= list_item_weight(v.len()) as i64;
396                        }
397                        l.drain(..s);
398                    }
399                }
400            }
401            d
402        };
403        self.account_delta(key, delta);
404        self.drop_if_empty_list(key);
405        Ok(())
406    }
407}