rustdis/
store.rs

1use bytes::Bytes;
2use std::collections::HashMap;
3use std::ops::AddAssign;
4use std::str::FromStr;
5
6pub struct Store {
7    store: HashMap<String, Bytes>,
8}
9
10impl Store {
11    pub fn new() -> Store {
12        Store {
13            store: HashMap::new(),
14        }
15    }
16
17    pub fn set(&mut self, key: String, value: Bytes) {
18        self.store.insert(key, value);
19    }
20
21    pub fn get(&self, key: &str) -> Option<&Bytes> {
22        self.store.get(key)
23    }
24
25    pub fn remove(&mut self, key: &str) -> Option<Bytes> {
26        self.store.remove(key)
27    }
28
29    pub fn exists(&self, key: &str) -> bool {
30        self.store.contains_key(key)
31    }
32
33    pub fn keys(&self) -> impl Iterator<Item = &String> {
34        self.store.keys()
35    }
36
37    pub fn size(&self) -> usize {
38        self.store.len()
39    }
40
41    pub fn iter(&self) -> impl Iterator<Item = (&String, &Bytes)> {
42        self.store.iter()
43    }
44
45    pub fn incr_by<T>(&mut self, key: &str, increment: T) -> Result<T, String>
46    where
47        T: FromStr + ToString + AddAssign + Default,
48    {
49        let err = "value is not of the correct type or out of range".to_string();
50
51        let mut value = match self.get(key) {
52            Some(value) => match std::str::from_utf8(value.as_ref())
53                .map_err(|_| err.clone())
54                .and_then(|s| s.parse::<T>().map_err(|_| err.clone()))
55            {
56                Ok(value) => value,
57                Err(e) => return Err(e),
58            },
59            None => T::default(),
60        };
61
62        value += increment;
63
64        self.set(key.to_string(), value.to_string().into());
65
66        Ok(value)
67    }
68}
69
70impl Default for Store {
71    fn default() -> Self {
72        Self::new()
73    }
74}