Skip to main content

kevy_store/
list.rs

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