perf: reduce heap alloc

crates/rle/src/lib.rs

@@ -32,5 +32,6 @@ pub use crate::rle_vec_old::{RleVecWithIndex, SearchResult, SliceIterator};
 pub mod rle_impl;
 pub use rle_tree::tree_trait::RleTreeTrait;
 pub use rle_tree::RleTree;
+mod small_set;
 #[cfg(test)]
 mod test;

crates/rle/src/rle_tree/node/internal_impl.rs

@@ -1,14 +1,17 @@
 use std::{
-    collections::{BinaryHeap, HashMap, HashSet},
+    collections::{BinaryHeap, HashSet},
     fmt::{Debug, Error, Formatter},
 };
 
-use fxhash::{FxBuildHasher, FxHashSet};
+use fxhash::FxHashSet;
 use smallvec::SmallVec;
 
-use crate::rle_tree::{
-    node::utils::distribute,
-    tree_trait::{FindPosResult, Position},
+use crate::{
+    rle_tree::{
+        node::utils::distribute,
+        tree_trait::{FindPosResult, Position},
+    },
+    small_set::SmallSet,
 };
 
 use super::*;
@@ -143,7 +146,7 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
         &mut self,
         from: Option<A::Int>,
         to: Option<A::Int>,
-        visited: &mut Vec<(usize, NonNull<Node<'a, T, A>>)>,
+        visited: &mut SmallVec<[(usize, NonNull<Node<'a, T, A>>); 8]>,
         depth: usize,
         notify: &mut F,
     ) -> Result<(), &'a mut Node<'a, T, A>>
@@ -523,7 +526,7 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
         F: FnMut(&T, *mut LeafNode<'_, T, A>),
     {
         debug_assert!(self.is_root());
-        let mut zipper = Vec::with_capacity(8);
+        let mut zipper = SmallVec::with_capacity(12);
         match self._delete(start, end, &mut zipper, 1, notify) {
             Ok(_) => {
                 A::update_cache_internal(self);
@@ -538,12 +541,9 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
         let removed = self._root_shrink_levels_if_one_child();
 
         // filter the same
-        let mut visited: HashSet<NonNull<_>, _> =
-            HashSet::with_capacity_and_hasher(zipper.len(), FxBuildHasher::default());
-        let mut should_skip: HashSet<NonNull<_>, _> =
-            HashSet::with_capacity_and_hasher(zipper.len(), FxBuildHasher::default());
-        let mut depth_to_node: HashMap<isize, SmallVec<[NonNull<_>; 2]>, _> =
-            HashMap::with_capacity_and_hasher(zipper.len(), FxBuildHasher::default());
+        let mut visited: SmallSet<NonNull<_>, 12> = SmallSet::new();
+        let mut should_skip: SmallSet<NonNull<_>, 12> = SmallSet::new();
+        let mut depth_to_node: SmallVec<[SmallVec<[NonNull<_>; 2]>; 10]> = smallvec::smallvec![];
         let mut zipper: BinaryHeap<(isize, NonNull<Node<'a, T, A>>)> = zipper
             .into_iter()
             .filter_map(|(i, mut ptr)| {
@@ -559,7 +559,10 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
                 if visited.contains(&ptr) {
                     None
                 } else {
-                    depth_to_node.entry(i as isize).or_default().push(ptr);
+                    while depth_to_node.len() <= i {
+                        depth_to_node.push(SmallVec::new());
+                    }
+                    depth_to_node[i].push(ptr);
                     visited.insert(ptr);
                     Some((-(i as isize), ptr))
                 }
@@ -600,7 +603,7 @@ impl<'a, T: Rle, A: RleTreeTrait<T>> InternalNode<'a, T, A> {
 
             if to_delete {
                 should_skip.insert(node_ptr);
-                if let Some(parent) = depth_to_node.get(&(-reverse_depth - 1)).as_ref() {
+                if let Some(parent) = depth_to_node.get((-reverse_depth - 1) as usize).as_ref() {
                     for ptr in parent.iter() {
                         zipper.push((reverse_depth + 1, *ptr));
                     }

crates/rle/src/small_set.rs

@@ -0,0 +1,95 @@
+use std::{hash::Hash, mem::MaybeUninit};
+
+use fxhash::FxHashSet;
+
+pub enum SmallSet<T, const SIZE: usize> {
+    Arr([Option<T>; SIZE]),
+    Set(FxHashSet<T>),
+}
+
+impl<T: Eq + Hash, const SIZE: usize> SmallSet<T, SIZE> {
+    pub fn new() -> Self {
+        let a: MaybeUninit<[Option<T>; SIZE]> = MaybeUninit::zeroed();
+        // SAFETY: we will init the array below
+        let mut a = unsafe { a.assume_init_read() };
+        for i in a.iter_mut() {
+            *i = None;
+        }
+        SmallSet::Arr(a)
+    }
+
+    /// Adds a value to the set.
+    ///
+    /// Returns whether the value was newly inserted. That is:
+    ///
+    /// - If the set did not previously contain this value, `true` is returned.
+    /// - If the set already contained this value, `false` is returned.
+    ///
+    pub fn insert(&mut self, v: T) -> bool {
+        match self {
+            SmallSet::Arr(a) => {
+                for i in a.iter_mut() {
+                    if let Some(i) = i {
+                        if *i == v {
+                            return false;
+                        }
+                    } else {
+                        *i = Some(v);
+                        return true;
+                    }
+                }
+
+                let mut set = FxHashSet::with_capacity_and_hasher(SIZE + 1, Default::default());
+                for i in a.iter_mut() {
+                    set.insert(std::mem::take(i).unwrap());
+                }
+
+                let ans = set.insert(v);
+                *self = SmallSet::Set(set);
+                ans
+            }
+            SmallSet::Set(set) => set.insert(v),
+        }
+    }
+
+    pub fn contains(&mut self, v: &T) -> bool {
+        match self {
+            SmallSet::Arr(a) => {
+                for i in a.iter_mut() {
+                    if let Some(i) = i {
+                        if i == v {
+                            return true;
+                        }
+                    } else {
+                        return false;
+                    }
+                }
+
+                false
+            }
+            SmallSet::Set(set) => set.contains(v),
+        }
+    }
+
+    pub fn remove(&mut self, v: &T) -> bool {
+        match self {
+            SmallSet::Arr(a) => {
+                for i in a.iter_mut() {
+                    if i.as_ref() == Some(v) {
+                        *i = None;
+                        return true;
+                    }
+                }
+
+                false
+            }
+            SmallSet::Set(set) => set.remove(v),
+        }
+    }
+}
+
+impl<T: Eq + Hash, const SIZE: usize> Default for SmallSet<T, SIZE> {
+    fn default() -> Self {
+        Self::new()
+    }
+}