autocommit 2022-01-30 22:41:40 CET

src/utils/smallmap.rs

@@ -0,0 +1,109 @@
+use crate::utils::clonecell::UnsafeCellCloneSafe;
+use crate::utils::ptr_ext::{MutPtrExt, PtrExt};
+use smallvec::SmallVec;
+use std::cell::UnsafeCell;
+use std::mem;
+
+pub struct SmallMap<K, V, const N: usize> {
+    m: UnsafeCell<SmallVec<[(K, V); N]>>,
+}
+
+impl<K, V, const N: usize> Default for SmallMap<K, V, N> {
+    fn default() -> Self {
+        Self {
+            m: Default::default()
+        }
+    }
+}
+
+impl<K: Eq, V, const N: usize> SmallMap<K, V, N> {
+    pub fn new() -> Self {
+        Self {
+            m: UnsafeCell::new(SmallVec::new_const()),
+        }
+    }
+
+    pub fn insert(&self, k: K, v: V) -> Option<V> {
+        unsafe {
+            let m = self.m.get().deref_mut();
+            for (ek, ev) in &mut *m {
+                if ek == &k {
+                    return Some(mem::replace(ev, v));
+                }
+            }
+            m.push((k, v));
+            None
+        }
+    }
+
+    pub fn remove(&self, k: &K) -> Option<V> {
+        unsafe {
+            let m = self.m.get().deref_mut();
+            for (idx, (ek, _)) in m.iter_mut().enumerate() {
+                if ek == k {
+                    return Some(m.swap_remove(idx).1);
+                }
+            }
+            None
+        }
+    }
+
+    pub fn take(&self) -> SmallVec<[(K, V); N]> {
+        unsafe {
+            mem::take(self.m.get().deref_mut())
+        }
+    }
+
+    pub fn pop(&self) -> Option<(K, V)> {
+        unsafe {
+            self.m.get().deref_mut().pop()
+        }
+    }
+}
+
+impl<K: Eq, V: UnsafeCellCloneSafe, const N: usize> SmallMap<K, V, N> {
+    pub fn get(&self, k: &K) -> Option<V> {
+        unsafe {
+            let m = self.m.get().deref();
+            for (ek, ev) in m {
+                if ek == k {
+                    return Some(ev.clone());
+                }
+            }
+            None
+        }
+    }
+}
+
+impl<'a, K: Copy, V: UnsafeCellCloneSafe, const N: usize> IntoIterator for &'a SmallMap<K, V, N> {
+    type Item = (K, V);
+    type IntoIter = SmallMapIter<'a, K, V, N>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SmallMapIter {
+            pos: 0,
+            map: self,
+        }
+    }
+}
+
+pub struct SmallMapIter<'a, K, V, const N: usize> {
+    pos: usize,
+    map: &'a SmallMap<K, V, N>,
+}
+
+impl<'a, K: Copy, V: UnsafeCellCloneSafe, const N: usize> Iterator for SmallMapIter<'a, K, V, N> {
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        unsafe {
+            let v = self.map.m.get().deref();
+            if self.pos >= v.len() {
+                return None;
+            }
+            let (k, v) = &v[self.pos];
+            self.pos += 1;
+            Some((*k, v.clone()))
+        }
+    }
+}