diff --git a/src/utils.rs b/src/utils.rs index 92928ca2..1ba53964 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -2,6 +2,7 @@ pub mod activation_token; pub mod array; pub mod array_to_tuple; pub mod asyncevent; +pub mod binary_search_map; pub mod bindings; pub mod bitfield; pub mod bitflags; diff --git a/src/utils/binary_search_map.rs b/src/utils/binary_search_map.rs new file mode 100644 index 00000000..0976e02b --- /dev/null +++ b/src/utils/binary_search_map.rs @@ -0,0 +1,215 @@ +use { + crate::utils::ptr_ext::{MutPtrExt, PtrExt}, + smallvec::SmallVec, + std::{ + fmt::{Debug, Formatter}, + mem, + }, +}; + +pub struct BinarySearchMap { + m: SmallVec<[(K, V); N]>, +} + +impl Debug for BinarySearchMap { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_map() + .entries(self.m.iter().map(|e| (&e.0, &e.1))) + .finish() + } +} + +impl Default for BinarySearchMap { + fn default() -> Self { + Self { + m: Default::default(), + } + } +} + +impl BinarySearchMap { + pub fn new_with(k: K, v: V) -> Self { + let mut sv = SmallVec::new(); + sv.push((k, v)); + Self { m: sv } + } + + pub fn new() -> Self { + Self { + m: SmallVec::new_const(), + } + } + + pub fn len(&self) -> usize { + self.m.len() + } + + fn pos(&self, k: &K) -> Result + where + K: Ord + Eq, + { + self.m.binary_search_by(|(c, _)| c.cmp(k)) + } + + pub fn contains(&self, k: &K) -> bool + where + K: Ord + Eq, + { + self.pos(k).is_ok() + } + + pub fn not_contains(&self, k: &K) -> bool + where + K: Ord + Eq, + { + !self.contains(k) + } + + pub fn insert(&mut self, k: K, v: V) -> Option + where + K: Ord + Eq, + { + match self.pos(&k) { + Ok(p) => Some(mem::replace(&mut self.m[p], (k, v)).1), + Err(p) => { + self.m.insert(p, (k, v)); + None + } + } + } + + pub fn get(&self, k: &K) -> Option<&V> + where + K: Ord + Eq, + { + self.pos(k).ok().map(|p| &self.m[p].1) + } + + pub fn get_mut(&mut self, k: &K) -> Option<&mut V> + where + K: Ord + Eq, + { + self.pos(k).ok().map(|p| &mut self.m[p].1) + } + + pub fn get_or_default_mut(&mut self, k: K) -> &mut V + where + K: Ord + Eq, + V: Default, + { + self.get_or_insert_with(k, || V::default()) + } + + pub fn get_or_insert_with(&mut self, k: K, f: F) -> &mut V + where + K: Ord + Eq, + F: FnOnce() -> V, + { + let p = match self.pos(&k) { + Ok(p) => return &mut self.m[p].1, + Err(p) => p, + }; + self.m.insert(p, (k, f())); + &mut self.m[p].1 + } + + pub fn is_empty(&self) -> bool { + self.m.is_empty() + } + + pub fn remove(&mut self, k: &K) -> Option + where + K: Ord + Eq, + { + if let Ok(p) = self.pos(k) { + return Some(self.m.remove(p).1); + } + None + } + + pub fn clear(&mut self) { + let _v = mem::replace(&mut self.m, SmallVec::new()); + } + + pub fn take(&mut self) -> SmallVec<[(K, V); N]> { + mem::take(&mut self.m) + } + + pub fn iter<'a>(&'a self) -> BinarySearchMapIter<'a, K, V, N> { + BinarySearchMapIter { pos: 0, map: self } + } + + pub fn values<'a>(&'a self) -> impl Iterator + 'a { + self.iter().map(|(_, v)| v) + } + + pub fn iter_mut<'a>(&'a mut self) -> BinarySearchMapMutIterMut<'a, K, V, N> { + BinarySearchMapMutIterMut { pos: 0, map: self } + } + + pub fn remove_if bool>(&mut self, mut f: F) { + let mut i = 0; + while i < self.m.len() { + let (k, v) = &self.m[i]; + if f(k, v) { + self.m.remove(i); + } else { + i += 1; + } + } + } +} + +impl<'a, K: Copy, V, const N: usize> IntoIterator for &'a BinarySearchMap { + type Item = (&'a K, &'a V); + type IntoIter = BinarySearchMapIter<'a, K, V, N>; + + fn into_iter(self) -> Self::IntoIter { + BinarySearchMapIter { pos: 0, map: self } + } +} + +impl<'a, K: Copy, V, const N: usize> IntoIterator for &'a mut BinarySearchMap { + type Item = (&'a K, &'a mut V); + type IntoIter = BinarySearchMapMutIterMut<'a, K, V, N>; + + fn into_iter(self) -> Self::IntoIter { + BinarySearchMapMutIterMut { pos: 0, map: self } + } +} + +pub struct BinarySearchMapIter<'a, K, V, const N: usize> { + pos: usize, + map: &'a BinarySearchMap, +} + +impl<'a, K, V, const N: usize> Iterator for BinarySearchMapIter<'a, K, V, N> { + type Item = (&'a K, &'a V); + + fn next(&mut self) -> Option { + if self.pos >= self.map.m.len() { + return None; + } + let (k, v) = &self.map.m[self.pos]; + self.pos += 1; + Some((k, v)) + } +} + +pub struct BinarySearchMapMutIterMut<'a, K, V, const N: usize> { + pos: usize, + map: &'a mut BinarySearchMap, +} + +impl<'a, K, V, const N: usize> Iterator for BinarySearchMapMutIterMut<'a, K, V, N> { + type Item = (&'a K, &'a mut V); + + fn next(&mut self) -> Option { + if self.pos >= self.map.m.len() { + return None; + } + let (k, v) = &mut self.map.m[self.pos]; + self.pos += 1; + unsafe { Some(((k as *const K).deref(), (v as *mut V).deref_mut())) } + } +}