ftl_jiter/
lazy_index_map.rs

1use std::borrow::{Borrow, Cow};
2use std::fmt;
3use std::hash::Hash;
4use std::slice::Iter as SliceIter;
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::OnceLock;
7
8use ahash::AHashMap;
9use smallvec::SmallVec;
10
11/// Like [IndexMap](https://docs.rs/indexmap/latest/indexmap/) but only builds the lookup map when it's needed.
12pub struct LazyIndexMap<K, V> {
13    vec: SmallVec<[(K, V); 8]>,
14    map: OnceLock<AHashMap<K, usize>>,
15    last_find: AtomicUsize,
16}
17
18impl<K, V> Default for LazyIndexMap<K, V>
19where
20    K: Clone + fmt::Debug + Eq + Hash,
21    V: fmt::Debug,
22{
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl<K: Clone, V: Clone> Clone for LazyIndexMap<K, V> {
29    fn clone(&self) -> Self {
30        Self {
31            vec: self.vec.clone(),
32            map: self.map.clone(),
33            last_find: AtomicUsize::new(0),
34        }
35    }
36}
37
38impl<K, V> fmt::Debug for LazyIndexMap<K, V>
39where
40    K: Clone + fmt::Debug + Eq + Hash,
41    V: fmt::Debug,
42{
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        f.debug_map().entries(self.iter_unique()).finish()
45    }
46}
47
48// picked to be a good tradeoff after experimenting with `lazy_map_lookup` benchmark, should cover most models
49const HASHMAP_THRESHOLD: usize = 16;
50
51/// Like [IndexMap](https://docs.rs/indexmap/latest/indexmap/) but only builds the lookup map when it's needed.
52impl<K, V> LazyIndexMap<K, V>
53where
54    K: Clone + fmt::Debug + Eq + Hash,
55    V: fmt::Debug,
56{
57    pub fn new() -> Self {
58        Self {
59            vec: SmallVec::new(),
60            map: OnceLock::new(),
61            last_find: AtomicUsize::new(0),
62        }
63    }
64
65    pub fn insert(&mut self, key: K, value: V) {
66        if let Some(map) = self.map.get_mut() {
67            map.insert(key.clone(), self.vec.len());
68        }
69        self.vec.push((key, value));
70    }
71
72    pub fn len(&self) -> usize {
73        self.get_map().len()
74    }
75
76    pub fn is_empty(&self) -> bool {
77        self.vec.is_empty()
78    }
79
80    pub fn get<Q>(&self, key: &Q) -> Option<&V>
81    where
82        K: Borrow<Q> + PartialEq<Q>,
83        Q: Hash + Eq + ?Sized,
84    {
85        let vec_len = self.vec.len();
86        // if the vec is longer than the threshold, we use the hashmap for lookups
87        if vec_len > HASHMAP_THRESHOLD {
88            self.get_map().get(key).map(|&i| &self.vec[i].1)
89        } else {
90            // otherwise we find the value in the vec
91            // we assume the most likely position for the match is at `last_find + 1`
92            let first_try = self.last_find.load(Ordering::Relaxed) + 1;
93            for i in first_try..first_try + vec_len {
94                let index = i % vec_len;
95                let (k, v) = &self.vec[index];
96                if k == key {
97                    self.last_find.store(index, Ordering::Relaxed);
98                    return Some(v);
99                }
100            }
101            None
102        }
103    }
104
105    pub fn keys(&self) -> impl Iterator<Item = &K> {
106        self.vec.iter().map(|(k, _)| k)
107    }
108
109    pub fn iter(&self) -> SliceIter<'_, (K, V)> {
110        self.vec.iter()
111    }
112
113    pub fn iter_unique(&self) -> impl Iterator<Item = (&K, &V)> {
114        IterUnique {
115            vec: &self.vec,
116            map: self.get_map(),
117            index: 0,
118        }
119    }
120
121    fn get_map(&self) -> &AHashMap<K, usize> {
122        self.map.get_or_init(|| {
123            self.vec
124                .iter()
125                .enumerate()
126                .map(|(index, (key, _))| (key.clone(), index))
127                .collect()
128        })
129    }
130}
131
132impl<'j> LazyIndexMap<Cow<'j, str>, crate::JsonValue<'j>> {
133    pub(crate) fn to_static(&self) -> LazyIndexMap<Cow<'static, str>, crate::JsonValue<'static>> {
134        LazyIndexMap {
135            vec: self
136                .vec
137                .iter()
138                .map(|(k, v)| (k.to_string().into(), v.to_static()))
139                .collect(),
140            map: OnceLock::new(),
141            last_find: AtomicUsize::new(0),
142        }
143    }
144}
145
146impl<K: PartialEq, V: PartialEq> PartialEq for LazyIndexMap<K, V> {
147    fn eq(&self, other: &Self) -> bool {
148        self.vec == other.vec
149    }
150}
151
152struct IterUnique<'a, K, V> {
153    vec: &'a SmallVec<[(K, V); 8]>,
154    map: &'a AHashMap<K, usize>,
155    index: usize,
156}
157
158impl<'a, K: Hash + Eq, V> Iterator for IterUnique<'a, K, V> {
159    type Item = (&'a K, &'a V);
160
161    fn next(&mut self) -> Option<Self::Item> {
162        while self.index < self.vec.len() {
163            let (k, v) = &self.vec[self.index];
164            if let Some(map_index) = self.map.get(k) {
165                if *map_index == self.index {
166                    self.index += 1;
167                    return Some((k, v));
168                }
169            }
170            self.index += 1;
171        }
172        None
173    }
174}
ftl_jiter/lazy_index_map.rs

ftl_jiter/
lazy_index_map.rs
