diff --git a/crates/rle/src/lib.rs b/crates/rle/src/lib.rs index c97c28a5..19a19a49 100644 --- a/crates/rle/src/lib.rs +++ b/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; diff --git a/crates/rle/src/rle_tree/node/internal_impl.rs b/crates/rle/src/rle_tree/node/internal_impl.rs index 958cecf3..11c2f4ec 100644 --- a/crates/rle/src/rle_tree/node/internal_impl.rs +++ b/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> InternalNode<'a, T, A> { &mut self, from: Option, to: Option, - visited: &mut Vec<(usize, NonNull>)>, + visited: &mut SmallVec<[(usize, NonNull>); 8]>, depth: usize, notify: &mut F, ) -> Result<(), &'a mut Node<'a, T, A>> @@ -523,7 +526,7 @@ impl<'a, T: Rle, A: RleTreeTrait> 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> InternalNode<'a, T, A> { let removed = self._root_shrink_levels_if_one_child(); // filter the same - let mut visited: HashSet, _> = - HashSet::with_capacity_and_hasher(zipper.len(), FxBuildHasher::default()); - let mut should_skip: HashSet, _> = - HashSet::with_capacity_and_hasher(zipper.len(), FxBuildHasher::default()); - let mut depth_to_node: HashMap; 2]>, _> = - HashMap::with_capacity_and_hasher(zipper.len(), FxBuildHasher::default()); + let mut visited: SmallSet, 12> = SmallSet::new(); + let mut should_skip: SmallSet, 12> = SmallSet::new(); + let mut depth_to_node: SmallVec<[SmallVec<[NonNull<_>; 2]>; 10]> = smallvec::smallvec![]; let mut zipper: BinaryHeap<(isize, NonNull>)> = zipper .into_iter() .filter_map(|(i, mut ptr)| { @@ -559,7 +559,10 @@ impl<'a, T: Rle, A: RleTreeTrait> 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> 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)); } diff --git a/crates/rle/src/small_set.rs b/crates/rle/src/small_set.rs new file mode 100644 index 00000000..89e79eea --- /dev/null +++ b/crates/rle/src/small_set.rs @@ -0,0 +1,95 @@ +use std::{hash::Hash, mem::MaybeUninit}; + +use fxhash::FxHashSet; + +pub enum SmallSet { + Arr([Option; SIZE]), + Set(FxHashSet), +} + +impl SmallSet { + pub fn new() -> Self { + let a: MaybeUninit<[Option; 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 Default for SmallSet { + fn default() -> Self { + Self::new() + } +}