Skip to main content

kevy_store/
list.rs

1//! `Store` list commands.
2
3use crate::util::{norm_index, range_bounds};
4use crate::value::{ListData, SmallBytes, Value, list_item_weight};
5use crate::{Entry, Store, StoreError};
6use std::sync::Arc;
7
8impl Store {
9    // ---- lists ---------------------------------------------------------
10
11    fn list_mut(&mut self, key: &[u8], create: bool) -> Result<Option<&mut ListData>, StoreError> {
12        if self.live_entry_mut(key).is_none() {
13            if !create {
14                return Ok(None);
15            }
16            self.insert_entry(
17                SmallBytes::from_slice(key),
18                Entry::new(Value::List(Arc::default()), None),
19            );
20        }
21        match &mut self.map.get_mut(key).expect("present").value {
22            Value::List(l) => Ok(Some(Arc::make_mut(l))),
23            _ => Err(StoreError::WrongType),
24        }
25    }
26
27    fn list_ref(&mut self, key: &[u8]) -> Result<Option<&ListData>, StoreError> {
28        match self.live_entry(key) {
29            None => Ok(None),
30            Some(e) => match &e.value {
31                Value::List(l) => Ok(Some(l.as_ref())),
32                _ => Err(StoreError::WrongType),
33            },
34        }
35    }
36
37    /// Remove `key` if it now holds an empty list.
38    fn drop_if_empty_list(&mut self, key: &[u8]) {
39        let empty = matches!(self.map.get(key).map(|e| &e.value), Some(Value::List(l)) if l.is_empty());
40        if empty {
41            self.remove_entry(key);
42        }
43    }
44
45    /// `LPUSH` — prepend each value in turn; returns the new length.
46    pub fn lpush(&mut self, key: &[u8], values: &[Vec<u8>]) -> Result<usize, StoreError> {
47        let (new_len, delta) = {
48            let l = self.list_mut(key, true)?.expect("created");
49            let mut d: i64 = 0;
50            for v in values {
51                d += list_item_weight(v.len()) as i64;
52                l.push_front(v.clone());
53            }
54            (l.len(), d)
55        };
56        self.account_delta(key, delta);
57        Ok(new_len)
58    }
59
60    /// `RPUSH` — append each value; returns the new length.
61    pub fn rpush(&mut self, key: &[u8], values: &[Vec<u8>]) -> Result<usize, StoreError> {
62        let (new_len, delta) = {
63            let l = self.list_mut(key, true)?.expect("created");
64            let mut d: i64 = 0;
65            for v in values {
66                d += list_item_weight(v.len()) as i64;
67                l.push_back(v.clone());
68            }
69            (l.len(), d)
70        };
71        self.account_delta(key, delta);
72        Ok(new_len)
73    }
74
75    /// `LPOP` — pop up to `count` from the head (deleting an emptied key).
76    pub fn lpop(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
77        let (out, delta) = {
78            let mut o = Vec::new();
79            let mut d: i64 = 0;
80            if let Some(l) = self.list_mut(key, false)? {
81                for _ in 0..count {
82                    match l.pop_front() {
83                        Some(v) => {
84                            d -= list_item_weight(v.len()) as i64;
85                            o.push(v);
86                        }
87                        None => break,
88                    }
89                }
90            }
91            (o, d)
92        };
93        self.account_delta(key, delta);
94        self.drop_if_empty_list(key);
95        Ok(out)
96    }
97
98    /// `RPOP` — pop up to `count` from the tail.
99    pub fn rpop(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
100        let (out, delta) = {
101            let mut o = Vec::new();
102            let mut d: i64 = 0;
103            if let Some(l) = self.list_mut(key, false)? {
104                for _ in 0..count {
105                    match l.pop_back() {
106                        Some(v) => {
107                            d -= list_item_weight(v.len()) as i64;
108                            o.push(v);
109                        }
110                        None => break,
111                    }
112                }
113            }
114            (o, d)
115        };
116        self.account_delta(key, delta);
117        self.drop_if_empty_list(key);
118        Ok(out)
119    }
120
121    pub fn llen(&mut self, key: &[u8]) -> Result<usize, StoreError> {
122        Ok(self.list_ref(key)?.map_or(0, std::collections::VecDeque::len))
123    }
124
125    pub fn lindex(&mut self, key: &[u8], idx: i64) -> Result<Option<Vec<u8>>, StoreError> {
126        match self.list_ref(key)? {
127            None => Ok(None),
128            Some(l) => Ok(norm_index(idx, l.len()).and_then(|i| l.get(i).cloned())),
129        }
130    }
131
132    pub fn lrange(
133        &mut self,
134        key: &[u8],
135        start: i64,
136        stop: i64,
137    ) -> Result<Vec<Vec<u8>>, StoreError> {
138        match self.list_ref(key)? {
139            None => Ok(Vec::new()),
140            Some(l) => Ok(match range_bounds(start, stop, l.len()) {
141                None => Vec::new(),
142                Some((s, e)) => l.iter().skip(s).take(e - s + 1).cloned().collect(),
143            }),
144        }
145    }
146
147    /// `LSET` — errors with `NoSuchKey` / `OutOfRange` like Redis.
148    pub fn lset(&mut self, key: &[u8], idx: i64, val: &[u8]) -> Result<(), StoreError> {
149        let delta = {
150            let l = self.list_mut(key, false)?.ok_or(StoreError::NoSuchKey)?;
151            let i = norm_index(idx, l.len()).ok_or(StoreError::OutOfRange)?;
152            let old_len = l[i].len() as i64;
153            l[i] = val.to_vec();
154            val.len() as i64 - old_len
155        };
156        self.account_delta(key, delta);
157        Ok(())
158    }
159
160    /// `LREM` — remove `count` occurrences of `val` (>0 head, <0 tail, 0 all).
161    pub fn lrem(&mut self, key: &[u8], count: i64, val: &[u8]) -> Result<usize, StoreError> {
162        let (removed, delta) = {
163            let mut r = 0usize;
164            let mut d: i64 = 0;
165            match self.list_mut(key, false)? {
166                None => (0, 0),
167                Some(l) => {
168                    if count >= 0 {
169                        let limit = if count == 0 {
170                            usize::MAX
171                        } else {
172                            count as usize
173                        };
174                        let mut i = 0;
175                        while i < l.len() {
176                            if r < limit && l[i] == val {
177                                d -= list_item_weight(l[i].len()) as i64;
178                                l.remove(i);
179                                r += 1;
180                            } else {
181                                i += 1;
182                            }
183                        }
184                    } else {
185                        let limit = (-count) as usize;
186                        let mut i = l.len();
187                        while i > 0 {
188                            i -= 1;
189                            if r < limit && l[i] == val {
190                                d -= list_item_weight(l[i].len()) as i64;
191                                l.remove(i);
192                                r += 1;
193                            }
194                        }
195                    }
196                    (r, d)
197                }
198            }
199        };
200        self.account_delta(key, delta);
201        self.drop_if_empty_list(key);
202        Ok(removed)
203    }
204
205    /// `LTRIM` — keep only `[start, stop]` (deleting an emptied key).
206    pub fn ltrim(&mut self, key: &[u8], start: i64, stop: i64) -> Result<(), StoreError> {
207        let delta = {
208            let mut d: i64 = 0;
209            if let Some(l) = self.list_mut(key, false)? {
210                match range_bounds(start, stop, l.len()) {
211                    None => {
212                        for v in l.iter() {
213                            d -= list_item_weight(v.len()) as i64;
214                        }
215                        l.clear();
216                    }
217                    Some((s, e)) => {
218                        for v in l.iter().skip(e + 1) {
219                            d -= list_item_weight(v.len()) as i64;
220                        }
221                        l.drain(e + 1..);
222                        for v in l.iter().take(s) {
223                            d -= list_item_weight(v.len()) as i64;
224                        }
225                        l.drain(..s);
226                    }
227                }
228            }
229            d
230        };
231        self.account_delta(key, delta);
232        self.drop_if_empty_list(key);
233        Ok(())
234    }
235}