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    /// `LINSERT key BEFORE|AFTER pivot value` — insert `value`
331    /// before/after the first occurrence of `pivot` in the list at
332    /// `key`. Returns:
333    /// - new list length on success (`>= 1`);
334    /// - `0` when `key` does not exist;
335    /// - `-1` when `pivot` was not found in the list.
336    ///
337    /// Matches Redis semantics.
338    pub fn linsert(
339        &mut self,
340        key: &[u8],
341        before: bool,
342        pivot: &[u8],
343        val: &[u8],
344    ) -> Result<i64, StoreError> {
345        self.promote_list_inline_to_heap(key);
346        let (result, delta) = match self.list_mut(key, false)? {
347            None => return Ok(0),
348            Some(l) => {
349                let Some(idx) = l.iter().position(|v| v.as_slice() == pivot) else {
350                    return Ok(-1);
351                };
352                let insert_at = if before { idx } else { idx + 1 };
353                l.insert(insert_at, val.to_vec());
354                let new_len = l.len();
355                let d = list_item_weight(val.len()) as i64;
356                (new_len as i64, d)
357            }
358        };
359        self.account_delta(key, delta);
360        Ok(result)
361    }
362
363    /// `LREM` — remove `count` occurrences (>0 head, <0 tail, 0 all).
364    pub fn lrem(&mut self, key: &[u8], count: i64, val: &[u8]) -> Result<usize, StoreError> {
365        self.promote_list_inline_to_heap(key);
366        let (removed, delta) = {
367            let mut r = 0usize;
368            let mut d: i64 = 0;
369            match self.list_mut(key, false)? {
370                None => (0, 0),
371                Some(l) => {
372                    if count >= 0 {
373                        let limit = if count == 0 {
374                            usize::MAX
375                        } else {
376                            count as usize
377                        };
378                        let mut i = 0;
379                        while i < l.len() {
380                            if r < limit && l[i] == val {
381                                d -= list_item_weight(l[i].len()) as i64;
382                                l.remove(i);
383                                r += 1;
384                            } else {
385                                i += 1;
386                            }
387                        }
388                    } else {
389                        let limit = (-count) as usize;
390                        let mut i = l.len();
391                        while i > 0 {
392                            i -= 1;
393                            if r < limit && l[i] == val {
394                                d -= list_item_weight(l[i].len()) as i64;
395                                l.remove(i);
396                                r += 1;
397                            }
398                        }
399                    }
400                    (r, d)
401                }
402            }
403        };
404        self.account_delta(key, delta);
405        self.drop_if_empty_list(key);
406        Ok(removed)
407    }
408
409    /// `LTRIM` — keep only `[start, stop]` (deleting emptied key).
410    pub fn ltrim(&mut self, key: &[u8], start: i64, stop: i64) -> Result<(), StoreError> {
411        self.promote_list_inline_to_heap(key);
412        let delta = {
413            let mut d: i64 = 0;
414            if let Some(l) = self.list_mut(key, false)? {
415                match range_bounds(start, stop, l.len()) {
416                    None => {
417                        for v in l.iter() {
418                            d -= list_item_weight(v.len()) as i64;
419                        }
420                        l.clear();
421                    }
422                    Some((s, e)) => {
423                        for v in l.iter().skip(e + 1) {
424                            d -= list_item_weight(v.len()) as i64;
425                        }
426                        l.drain(e + 1..);
427                        for v in l.iter().take(s) {
428                            d -= list_item_weight(v.len()) as i64;
429                        }
430                        l.drain(..s);
431                    }
432                }
433            }
434            d
435        };
436        self.account_delta(key, delta);
437        self.drop_if_empty_list(key);
438        Ok(())
439    }
440}