vsdb/basic_multi_key/mapx_raw/
mod.rs

1//!
2//! A multi-key version of `MapxRaw`.
3//!
4//! `MapxRawMk` (MapxRaw Multi-Key) provides a map-like interface where each value
5//! is associated with a sequence of keys. This allows for creating nested or
6//! hierarchical key structures with raw byte keys and values.
7//!
8//! # Examples
9//!
10//! ```
11//! use vsdb::basic_multi_key::mapx_raw::MapxRawMk;
12//! use vsdb::{vsdb_set_base_dir, vsdb_get_base_dir};
13//! use std::fs;
14//!
15//! // It's recommended to use a temporary directory for testing
16//! let dir = format!("/tmp/vsdb_testing/{}", rand::random::<u128>());
17//! vsdb_set_base_dir(&dir).unwrap();
18//!
19//! let mut m = MapxRawMk::new(2); // Two keys
20//!
21//! // Insert a value with a multi-key
22//! m.insert(&[&[1], &[10]], &[100]).unwrap();
23//! m.insert(&[&[2], &[20]], &[200]).unwrap();
24//!
25//! // Get a value
26//! assert_eq!(m.get(&[&[1], &[10]]), Some(vec![100]));
27//!
28//! // Remove a value
29//! m.remove(&[&[1], &[10]]).unwrap();
30//! assert!(!m.contains_key(&[&[1], &[10]]));
31//!
32//! // Clean up the directory
33//! fs::remove_dir_all(vsdb_get_base_dir()).unwrap();
34//! ```
35
36#[cfg(test)]
37mod test;
38
39use crate::{
40    MapxRaw,
41    common::{RawKey, RawValue, ende::ValueEnDe},
42};
43use ruc::*;
44use serde::{Deserialize, Serialize};
45use std::ops::{Deref, DerefMut};
46
47/// A multi-key map with raw keys and values.
48#[derive(Clone, Serialize, Deserialize, Debug)]
49#[serde(bound = "")]
50pub struct MapxRawMk {
51    // Will never be changed once created
52    key_size: u32,
53    // A nested map-structure, looks like:
54    // map { key => map { key => map { key => value } } }
55    inner: MapxRaw,
56}
57
58impl MapxRawMk {
59    /// Creates a "shadow" copy of the `MapxRawMk` instance.
60    ///
61    /// # Safety
62    ///
63    /// This API breaks Rust's semantic safety guarantees. Use only in a race-free environment.
64    #[inline(always)]
65    pub unsafe fn shadow(&self) -> Self {
66        unsafe {
67            Self {
68                key_size: self.key_size,
69                inner: self.inner.shadow(),
70            }
71        }
72    }
73
74    /// Creates a new `MapxRawMk` with a specified number of keys.
75    ///
76    /// # Panics
77    ///
78    /// Panics if `key_size` is 0.
79    #[inline(always)]
80    pub fn new(key_size: u32) -> Self {
81        assert!(0 < key_size);
82        Self {
83            key_size,
84            inner: MapxRaw::new(),
85        }
86    }
87
88    /// Retrieves a value from the map for a given multi-key.
89    #[inline(always)]
90    pub fn get(&self, key: &[&[u8]]) -> Option<RawValue> {
91        if key.len() != self.key_size as usize {
92            return None;
93        }
94
95        let mut hdr = unsafe { self.inner.shadow() };
96        for (idx, k) in key.iter().enumerate() {
97            if let Some(v) = hdr.get(k) {
98                if 1 + idx == self.key_size as usize {
99                    return Some(v);
100                } else {
101                    hdr = pnk!(ValueEnDe::decode(&v));
102                }
103            } else {
104                return None;
105            }
106        }
107
108        None // empty key
109    }
110
111    /// Retrieves a mutable reference to a value in the map.
112    #[inline(always)]
113    pub fn get_mut<'a>(&'a mut self, key: &'a [&'a [u8]]) -> Option<ValueMut<'a>> {
114        self.get(key).map(move |v| ValueMut::new(self, key, v))
115    }
116
117    /// Mocks a mutable value for a given key.
118    #[inline(always)]
119    pub(crate) fn mock_value_mut<'a>(
120        &'a mut self,
121        key: &'a [&'a [u8]],
122        v: RawValue,
123    ) -> ValueMut<'a> {
124        ValueMut::new(self, key, v)
125    }
126
127    /// Checks if the map contains a value for the specified multi-key.
128    #[inline(always)]
129    pub fn contains_key(&self, key: &[&[u8]]) -> bool {
130        self.get(key).is_some()
131    }
132
133    /// Checks if the map is empty.
134    #[inline(always)]
135    pub fn is_empty(&self) -> bool {
136        self.inner.is_empty()
137    }
138
139    /// Gets an entry for a given key, allowing for in-place modification.
140    #[inline(always)]
141    pub fn entry<'a>(&'a mut self, key: &'a [&'a [u8]]) -> Result<Entry<'a>> {
142        if key.len() != self.key_size() as usize {
143            Err(eg!())
144        } else {
145            Ok(Entry { key, hdr: self })
146        }
147    }
148
149    /// Inserts a key-value pair into the map.
150    #[inline(always)]
151    pub fn insert(&mut self, key: &[&[u8]], value: &[u8]) -> Result<Option<RawValue>> {
152        if key.len() != self.key_size as usize {
153            return Err(eg!("Incorrect key size"));
154        }
155
156        let mut ret = None;
157
158        let mut hdr = unsafe { self.inner.shadow() };
159        for (idx, k) in key.iter().enumerate() {
160            if 1 + idx == self.key_size as usize {
161                ret = hdr.insert(k, value);
162                break;
163            } else {
164                let mut new_hdr = None;
165                let f = || {
166                    new_hdr.replace(MapxRaw::new());
167                    new_hdr.as_ref().unwrap().encode()
168                };
169                let mutv = hdr.entry(k).or_insert_with(f);
170                let h = if let Some(h) = new_hdr {
171                    h
172                } else {
173                    pnk!(ValueEnDe::decode(mutv.as_ref()))
174                };
175                drop(mutv);
176                hdr = h;
177            }
178        }
179
180        Ok(ret)
181    }
182
183    /// Removes a key-value pair from the map. Supports batch removal by providing a partial key.
184    #[inline(always)]
185    pub fn remove(&mut self, key: &[&[u8]]) -> Result<Option<RawValue>> {
186        // Support batch removal from key path.
187        if key.len() > self.key_size as usize {
188            return Err(eg!("Incorrect key size"));
189        }
190
191        let mut hdr = unsafe { self.inner.shadow() };
192        for (idx, k) in key.iter().enumerate() {
193            if let Some(v) = hdr.get(k) {
194                // NOTE: use `key.len()` instead of `self.key_size`
195                if 1 + idx == key.len() {
196                    let ret = hdr.remove(k);
197                    // NOTE: use `self.key_size` instead of `key.len()`
198                    if 1 + idx == self.key_size as usize {
199                        return Ok(ret);
200                    } else {
201                        return Ok(None);
202                    }
203                } else {
204                    hdr = pnk!(ValueEnDe::decode(&v));
205                }
206            } else {
207                return Ok(None);
208            }
209        }
210
211        Ok(None) // empty key
212    }
213
214    /// Clears the map, removing all key-value pairs.
215    #[inline(always)]
216    pub fn clear(&mut self) {
217        self.inner.clear();
218    }
219
220    /// Checks if this `MapxRawMk` instance is the same as another.
221    #[inline(always)]
222    pub fn is_the_same_instance(&self, other_hdr: &Self) -> bool {
223        self.inner.is_the_same_instance(&other_hdr.inner)
224    }
225
226    /// Returns the number of keys in the map.
227    #[inline(always)]
228    pub fn key_size(&self) -> u32 {
229        self.key_size
230    }
231
232    /// Iterates over the map's entries, applying a function to each.
233    #[inline(always)]
234    pub fn iter_op<F>(&self, op: &mut F) -> Result<()>
235    where
236        F: FnMut(&[&[u8]], &[u8]) -> Result<()>,
237    {
238        self.iter_op_with_key_prefix(op, &[]).c(d!())
239    }
240
241    /// Iterates over the map's entries with a given key prefix, applying a function to each.
242    #[inline(always)]
243    pub fn iter_op_with_key_prefix<F>(
244        &self,
245        op: &mut F,
246        key_prefix: &[&[u8]],
247    ) -> Result<()>
248    where
249        F: FnMut(&[&[u8]], &[u8]) -> Result<()>,
250    {
251        let key_size = self.key_size() as usize;
252        let mut key_buf = vec![RawKey::default(); key_size];
253        let mut hdr = unsafe { self.inner.shadow() };
254        let mut depth = key_size;
255
256        if key_size < key_prefix.len() {
257            return Err(eg!("Invalid key size"));
258        } else {
259            for (idx, k) in key_prefix.iter().enumerate() {
260                if let Some(v) = hdr.get(k) {
261                    key_buf[idx] = k.to_vec();
262                    if 1 + idx == key_size {
263                        let key = key_buf
264                            .iter()
265                            .map(|sub_k| sub_k.as_ref())
266                            .collect::<Vec<_>>();
267                        return op(key.as_slice(), &v).c(d!());
268                    } else {
269                        hdr = pnk!(ValueEnDe::decode(&v));
270                        depth -= 1;
271                    }
272                } else {
273                    // key-prefix does not exist
274                    return Ok(());
275                }
276            }
277        };
278
279        self.recursive_walk(hdr, key_buf.as_mut_slice(), depth as u32, op)
280            .c(d!())
281    }
282
283    fn recursive_walk<F>(
284        &self,
285        hdr: MapxRaw,
286        key_buf: &mut [RawKey],
287        depth: u32,
288        op: &mut F,
289    ) -> Result<()>
290    where
291        F: FnMut(&[&[u8]], &[u8]) -> Result<()>,
292    {
293        let idx = (self.key_size() - depth) as usize;
294        if 1 == depth {
295            for (k, v) in hdr.iter() {
296                key_buf[idx] = k;
297                let key = key_buf
298                    .iter()
299                    .map(|sub_k| sub_k.as_ref())
300                    .collect::<Vec<_>>();
301                op(key.as_slice(), &v[..]).c(d!())?;
302            }
303        } else {
304            for (k, v) in hdr.iter() {
305                key_buf[idx] = k;
306                let hdr = pnk!(ValueEnDe::decode(&v));
307                self.recursive_walk(hdr, key_buf, depth - 1, op).c(d!())?;
308            }
309        }
310
311        Ok(())
312    }
313
314    #[inline(always)]
315    pub(super) fn iter_op_typed_value<V, F>(&self, op: &mut F) -> Result<()>
316    where
317        F: FnMut(&[&[u8]], &V) -> Result<()>,
318        V: ValueEnDe,
319    {
320        self.iter_op_typed_value_with_key_prefix(op, &[]).c(d!())
321    }
322
323    /// Iterates over the map's entries with a given key prefix, applying a function to each typed value.
324    #[inline(always)]
325    pub fn iter_op_typed_value_with_key_prefix<V, F>(
326        &self,
327        op: &mut F,
328        key_prefix: &[&[u8]],
329    ) -> Result<()>
330    where
331        F: FnMut(&[&[u8]], &V) -> Result<()>,
332        V: ValueEnDe,
333    {
334        let key_size = self.key_size() as usize;
335        let mut key_buf = vec![RawKey::default(); key_size];
336        let mut hdr = unsafe { self.inner.shadow() };
337        let mut depth = key_size;
338
339        if key_size < key_prefix.len() {
340            return Err(eg!("Invalid key size"));
341        } else {
342            for (idx, k) in key_prefix.iter().enumerate() {
343                if let Some(v) = hdr.get(k) {
344                    key_buf[idx] = k.to_vec();
345                    if 1 + idx == key_size {
346                        let key = key_buf
347                            .iter()
348                            .map(|sub_k| sub_k.as_ref())
349                            .collect::<Vec<_>>();
350                        return op(key.as_slice(), &pnk!(ValueEnDe::decode(&v))).c(d!());
351                    } else {
352                        hdr = pnk!(ValueEnDe::decode(&v));
353                        depth -= 1;
354                    }
355                } else {
356                    // key-prefix does not exist
357                    return Ok(());
358                }
359            }
360        };
361
362        self.recursive_walk_typed_value(hdr, key_buf.as_mut_slice(), depth as u32, op)
363            .c(d!())
364    }
365
366    fn recursive_walk_typed_value<V, F>(
367        &self,
368        hdr: MapxRaw,
369        key_buf: &mut [RawKey],
370        depth: u32,
371        op: &mut F,
372    ) -> Result<()>
373    where
374        F: FnMut(&[&[u8]], &V) -> Result<()>,
375        V: ValueEnDe,
376    {
377        let idx = (self.key_size() - depth) as usize;
378        if 1 == depth {
379            for (k, v) in hdr.iter() {
380                key_buf[idx] = k;
381                let key = key_buf
382                    .iter()
383                    .map(|sub_k| sub_k.as_ref())
384                    .collect::<Vec<_>>();
385                op(key.as_slice(), &pnk!(ValueEnDe::decode(&v))).c(d!())?;
386            }
387        } else {
388            for (k, v) in hdr.iter() {
389                key_buf[idx] = k;
390                let hdr = pnk!(ValueEnDe::decode(&v));
391                self.recursive_walk_typed_value(hdr, key_buf, depth - 1, op)
392                    .c(d!())?;
393            }
394        }
395
396        Ok(())
397    }
398
399    // TODO
400    // pub fn iter_mut_op
401    // pub fn iter_mut_op_with_key_prefix
402    // pub fn iter_mut_op_typed_value
403    // pub fn iter_mut_op_typed_value_with_key_prefix
404}
405
406/// A mutable reference to a value in a `MapxRawMk`.
407#[derive(Debug)]
408pub struct ValueMut<'a> {
409    hdr: &'a mut MapxRawMk,
410    key: &'a [&'a [u8]],
411    value: RawValue,
412}
413
414impl<'a> ValueMut<'a> {
415    fn new(hdr: &'a mut MapxRawMk, key: &'a [&'a [u8]], value: RawValue) -> Self {
416        ValueMut { hdr, key, value }
417    }
418}
419
420impl Drop for ValueMut<'_> {
421    fn drop(&mut self) {
422        pnk!(self.hdr.insert(self.key, &self.value));
423    }
424}
425
426impl Deref for ValueMut<'_> {
427    type Target = RawValue;
428    fn deref(&self) -> &Self::Target {
429        &self.value
430    }
431}
432
433impl DerefMut for ValueMut<'_> {
434    fn deref_mut(&mut self) -> &mut Self::Target {
435        &mut self.value
436    }
437}
438
439/// A view into a single entry in a map, which may either be vacant or occupied.
440pub struct Entry<'a> {
441    key: &'a [&'a [u8]],
442    hdr: &'a mut MapxRawMk,
443}
444
445impl<'a> Entry<'a> {
446    /// Ensures a value is in the entry by inserting the default if empty,
447    /// and returns a mutable reference to the value.
448    pub fn or_insert(self, default: &'a [u8]) -> ValueMut<'a> {
449        let hdr = self.hdr as *mut MapxRawMk;
450        match unsafe { &mut *hdr }.get_mut(self.key) {
451            Some(v) => v,
452            _ => unsafe { &mut *hdr }.mock_value_mut(self.key, default.to_vec()),
453        }
454    }
455}
456
457/////////////////////////////////////////////////////////////////////////////
458/////////////////////////////////////////////////////////////////////////////
vsdb/basic_multi_key/mapx_raw/mod.rs

vsdb/basic_multi_key/mapx_raw/
mod.rs
