pink_kv_session/
trackers.rs

1use alloc::vec;
2use alloc::{collections::BTreeSet, vec::Vec};
3
4use crate::traits::{AccessTracking, Key};
5
6pub type ReadTracker = AccessTracker<false>;
7pub type RwTracker = AccessTracker<true>;
8
9/// Tracker that always emits the given one access history
10pub struct OneLock {
11    lock_key: Key,
12    write_lock: bool,
13}
14
15impl OneLock {
16    pub fn new(lock_key: &[u8], write_lock: bool) -> Self {
17        Self {
18            lock_key: lock_key.to_vec(),
19            write_lock,
20        }
21    }
22}
23
24impl AccessTracking for OneLock {
25    // One key tracker tracks nothing
26    fn read(&mut self, _key: &[u8]) {}
27
28    fn write(&mut self, _key: &[u8]) {}
29
30    fn collect_into(self) -> (Vec<Key>, Vec<Key>) {
31        (
32            vec![self.lock_key.clone()],
33            if self.write_lock {
34                vec![self.lock_key]
35            } else {
36                vec![]
37            },
38        )
39    }
40}
41
42/// Tracker that emit read and(optional) write access history
43pub struct AccessTracker<const TRACK_WRITE: bool> {
44    track: BTreeSet<Key>,
45    writes: BTreeSet<Key>,
46}
47
48impl<const TRACK_WRITE: bool> AccessTracker<TRACK_WRITE> {
49    #[allow(clippy::new_without_default)]
50    pub fn new() -> Self {
51        Self {
52            track: Default::default(),
53            writes: Default::default(),
54        }
55    }
56}
57
58impl<const TRACK_WRITE: bool> AccessTracking for AccessTracker<TRACK_WRITE> {
59    fn read(&mut self, key: &[u8]) {
60        self.track.insert(key.to_vec());
61    }
62
63    fn write(&mut self, key: &[u8]) {
64        self.writes.insert(key.to_vec());
65        if TRACK_WRITE {
66            self.track.insert(key.to_vec());
67        }
68    }
69
70    fn collect_into(self) -> (Vec<Key>, Vec<Key>) {
71        (
72            self.track.into_iter().collect(),
73            self.writes.into_iter().collect(),
74        )
75    }
76}
pink_kv_session/trackers.rs

pink_kv_session/
trackers.rs
