onion_vm/utils/
fastmap.rs

1use std::{hash::Hash, sync::Arc};
2
3use rustc_hash::FxHashMap;
4
5#[derive(Debug, Clone)]
6pub struct OnionKeyPool<K: PartialEq + Eq + Hash + Clone> {
7    keys: Arc<[K]>,
8    key_to_index: Arc<FxHashMap<K, usize>>,
9}
10
11impl<K: PartialEq + Eq + Hash + Clone> OnionKeyPool<K> {
12    pub fn create(keys: Vec<K>) -> Self {
13        let mut key_to_index = FxHashMap::default();
14        for (i, k) in keys.iter().enumerate() {
15            key_to_index.insert(k.clone(), i);
16        }
17        Self {
18            keys: keys.into(),
19            key_to_index: Arc::new(key_to_index),
20        }
21    }
22
23    pub fn new(k: Arc<[K]>, i: Arc<FxHashMap<K, usize>>) -> Self {
24        Self {
25            keys: k,
26            key_to_index: i,
27        }
28    }
29    pub fn keys(&self) -> &[K] {
30        &self.keys
31    }
32
33    pub fn indices(&self) -> &FxHashMap<K, usize> {
34        &self.key_to_index
35    }
36}
37
38#[derive(Debug, Clone)]
39pub struct OnionFastMap<K: PartialEq + Eq + Hash + Clone, V> {
40    pairs: Vec<(usize, V)>,
41    pool: OnionKeyPool<K>,
42}
43
44impl<K: PartialEq + Eq + Hash + Clone, V> OnionFastMap<K, V> {
45    pub fn new(pool: OnionKeyPool<K>) -> Self {
46        Self {
47            pairs: Vec::new(),
48            pool,
49        }
50    }
51
52    pub fn pool(&self) -> &OnionKeyPool<K> {
53        &self.pool
54    }
55
56    pub fn new_with_pairs(pairs: Vec<(usize, V)>, pool: OnionKeyPool<K>) -> Self {
57        Self { pairs, pool }
58    }
59
60    #[inline(always)]
61    pub fn push<Q: ?Sized>(&mut self, key: &Q, value: V) -> Option<()>
62    where
63        K: std::borrow::Borrow<Q>,
64        Q: std::hash::Hash + Eq,
65    {
66        if let Some(index) = self.pool.indices().get(key).copied() {
67            self.pairs.push((index, value));
68            Some(())
69        } else {
70            None
71        }
72    }
73
74    #[inline(always)]
75    pub fn push_with_index(&mut self, index: usize, value: V) -> Option<()> {
76        if index < self.pool.keys().len() {
77            self.pairs.push((index, value));
78            Some(())
79        } else {
80            None
81        }
82    }
83
84    #[inline(always)]
85    pub fn pairs(&self) -> &[(usize, V)] {
86        &self.pairs
87    }
88
89    #[inline(always)]
90    pub fn set_pairs(&mut self, pairs: Vec<(usize, V)>) {
91        self.pairs = pairs;
92    }
93
94    #[inline(always)]
95    pub fn clear(&mut self) {
96        self.pairs.clear();
97    }
98
99    #[inline(always)]
100    pub fn to_index<Q: ?Sized>(&self, key: &Q) -> Option<usize>
101    where
102        K: std::borrow::Borrow<Q>,
103        Q: std::hash::Hash + Eq,
104    {
105        self.pool.indices().get(key).copied()
106    }
107
108    #[inline(always)]
109    pub fn from_index(&self, index: usize) -> Option<&K> {
110        self.pool.keys().get(index)
111    }
112
113    #[inline(always)]
114    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
115    where
116        K: std::borrow::Borrow<Q>,
117        Q: std::hash::Hash + Eq,
118    {
119        let target_id = self.to_index(key)?;
120        self.pairs
121            .iter()
122            .rfind(|(id, _)| *id == target_id)
123            .map(|(_, v)| v)
124    }
125    #[inline(always)]
126    pub fn get_by_index(&self, target_index: usize) -> Option<&V> {
127        // 同样，需要线性扫描
128        self.pairs
129            .iter()
130            .rfind(|(id, _)| *id == target_index)
131            .map(|(_, v)| v)
132    }
133}
134
135impl<K: PartialEq + Eq + Hash + Clone, V> Default for OnionFastMap<K, V> {
136    fn default() -> Self {
137        Self::new(OnionKeyPool::create(Vec::new()))
138    }
139}
onion_vm/utils/fastmap.rs

onion_vm/utils/
fastmap.rs
