dfir_rs/util/
monotonic_map.rs

1//! Module for [`MonotonicMap`].
2
3use super::clear::Clear;
4
5/// A map-like interface which in reality only stores one value at a time.
6///
7/// The keys must be monotonically increasing (i.e. timestamps). For DFIR, this allows state
8/// to be stored which resets each tick by using the tick counter as the key. In the generic `Map`
9/// case it can be swapped out for a true map to allow processing of multiple ticks of data at
10/// once.
11#[derive(Clone, Debug)]
12pub struct MonotonicMap<K, V>
13where
14    K: PartialOrd,
15{
16    key: Option<K>,
17    val: V,
18}
19
20impl<K, V> Default for MonotonicMap<K, V>
21where
22    K: PartialOrd,
23    V: Default,
24{
25    fn default() -> Self {
26        Self {
27            key: None,
28            val: Default::default(),
29        }
30    }
31}
32
33impl<K, V> MonotonicMap<K, V>
34where
35    K: PartialOrd,
36{
37    /// Creates a new `MonotonicMap` initialized with the given value. The
38    /// vaue will be `Clear`ed before it is accessed.
39    pub fn new_init(val: V) -> Self {
40        Self { key: None, val }
41    }
42
43    /// Inserts the value using the function if new `key` is strictly later than the current key.
44    pub fn get_mut_with(&mut self, key: K, init: impl FnOnce() -> V) -> &mut V {
45        if self.key.as_ref().is_none_or(|old_key| old_key < &key) {
46            self.key = Some(key);
47            self.val = (init)();
48        }
49        &mut self.val
50    }
51
52    // /// Returns the value for the monotonically increasing key, or `None` if
53    // /// the key has already passed.
54    // pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
55    //     if self
56    //         .key
57    //         .as_ref()
58    //         .map_or(true, |old_key| old_key <= &key)
59    //     {
60    //         self.key = Some(key);
61    //         Some(&mut self.val)
62    //     } else {
63    //         None
64    //     }
65    // }
66}
67
68impl<K, V> MonotonicMap<K, V>
69where
70    K: PartialOrd,
71    V: Default,
72{
73    /// Gets a mutable reference to the inner value. If `key` is strictly later than the existing
74    /// key, the value will be reset to `Default::default`.
75    pub fn get_mut_default(&mut self, key: K) -> &mut V {
76        self.get_mut_with(key, Default::default)
77    }
78}
79
80impl<K, V> MonotonicMap<K, V>
81where
82    K: PartialOrd,
83    V: Clear,
84{
85    /// Gets a mutable reference to the inner value. If `key` is strictly later than the existing
86    /// key, the value will be cleared via the [`Clear`] trait.
87    pub fn get_mut_clear(&mut self, key: K) -> &mut V {
88        if self.key.as_ref().is_none_or(|old_key| old_key < &key) {
89            self.key = Some(key);
90            self.val.clear();
91        }
92        &mut self.val
93    }
94}