diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 79edf65..a3657bc 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -38,6 +38,9 @@ jobs: - name: Cargo test run: cargo test --verbose + - name: Cargo test w/ litemap + run: cargo test --verbose --features litemap,litemap/testing + - name: Cargo test w/ serde run: cargo test --verbose --features serde @@ -50,7 +53,7 @@ jobs: - name: Cargo test all features if: matrix.toolchain == 'nightly' - run: cargo test --verbose --all-features + run: cargo test --verbose --all-features --features litemap/testing - name: Cargo bench if: matrix.toolchain == 'nightly' diff --git a/Cargo.toml b/Cargo.toml index d90e58f..fcbf1c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ extract_if = [] [dependencies] bytes = { version = "1", optional = true, default-features = false } +litemap = { version = "0.8", optional = true, default-features = false } serde = { version = "1", optional = true, default-features = false } malloc_size_of = { version = "0.1.1", optional = true, default-features = false } diff --git a/src/lib.rs b/src/lib.rs index 42a60ef..5f0ace2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,6 +30,12 @@ //! When this optional dependency is enabled, `SmallVec` implements the `serde::Serialize` and //! `serde::Deserialize` traits. //! +//! ### [`litemap`](https://docs.rs/litemap/0.8/litemap/) +//! +//! When this optional dependency is enabled, `SmallVec` implements the +//! [`store`](https://docs.rs/litemap/0.8/litemap/store/index.html) interface to enable pluggable +//! backends for LiteMap. +//! //! ### `extract_if` //! //! **This feature is unstable.** It may change to match the unstable `extract_if` method in libstd. @@ -94,6 +100,11 @@ use core::ptr::NonNull; #[cfg(feature = "bytes")] use bytes::{buf::UninitSlice, BufMut}; +#[cfg(feature = "litemap")] +use litemap::store::{ + Store, StoreBulkMut, StoreConstEmpty, StoreFromIterable, StoreFromIterator, StoreIntoIterator, + StoreIterable, StoreIterableMut, StoreMut, StoreSlice, +}; #[cfg(feature = "malloc_size_of")] use malloc_size_of::{MallocShallowSizeOf, MallocSizeOf, MallocSizeOfOps}; #[cfg(feature = "serde")] @@ -3022,3 +3033,269 @@ unsafe impl BufMut for SmallVec { self.resize(new_len, val); } } + +#[cfg(feature = "litemap")] +type MapF = fn(&(K, V)) -> (&K, &V); +#[cfg(feature = "litemap")] +#[inline] +fn map_f(input: &(K, V)) -> (&K, &V) { + (&input.0, &input.1) +} +#[cfg(feature = "litemap")] +type MapFMut = fn(&mut (K, V)) -> (&K, &mut V); +#[cfg(feature = "litemap")] +#[inline] +fn map_f_mut(input: &mut (K, V)) -> (&K, &mut V) { + (&input.0, &mut input.1) +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl Store for SmallVec<(K, V), N> { + #[inline] + fn lm_len(&self) -> usize { + self.len() + } + + #[inline] + fn lm_is_empty(&self) -> bool { + self.is_empty() + } + + #[inline] + fn lm_get(&self, index: usize) -> Option<(&K, &V)> { + self.get(index).map(map_f) + } + + #[inline] + fn lm_last(&self) -> Option<(&K, &V)> { + self.last().map(map_f) + } + + #[inline] + fn lm_binary_search_by(&self, mut cmp: F) -> Result + where + F: FnMut(&K) -> core::cmp::Ordering, + { + self.binary_search_by(|(k, _)| cmp(k)) + } +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreSlice for SmallVec<(K, V), N> { + type Slice = [(K, V)]; + + #[inline] + fn lm_get_range(&self, range: core::ops::Range) -> Option<&Self::Slice> { + self.get(range) + } +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreMut for SmallVec<(K, V), N> { + #[inline] + fn lm_with_capacity(capacity: usize) -> Self { + Self::with_capacity(capacity) + } + + #[inline] + fn lm_reserve(&mut self, additional: usize) { + self.reserve(additional) + } + + #[inline] + fn lm_get_mut(&mut self, index: usize) -> Option<(&K, &mut V)> { + self.as_mut_slice().get_mut(index).map(map_f_mut) + } + + #[inline] + fn lm_push(&mut self, key: K, value: V) { + self.push((key, value)) + } + + #[inline] + fn lm_insert(&mut self, index: usize, key: K, value: V) { + self.insert(index, (key, value)) + } + + #[inline] + fn lm_remove(&mut self, index: usize) -> (K, V) { + self.remove(index) + } + + #[inline] + fn lm_clear(&mut self) { + self.clear() + } +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreBulkMut for SmallVec<(K, V), N> { + #[inline] + fn lm_retain(&mut self, mut predicate: F) + where + F: FnMut(&K, &V) -> bool, + { + self.retain(|(k, v)| predicate(k, v)); + } + + /// This is copied from `litemap` crate, since this is not exported. + /// + /// Extends this store with items from an iterator. + /// + /// It uses a two-pass (sort + dedup) approach to avoid any potential quadratic costs. + /// + /// The asymptotic worst case complexity is O((n + m) log(n + m)), where `n` + /// is the number of elements already in `self` and `m` is the number of elements + /// in the iterator. The best case complexity is O(m), when the input iterator is + /// already sorted, keys aren't duplicated and all keys sort after the existing ones. + #[inline] + fn lm_extend(&mut self, other: I) + where + I: IntoIterator, + { + // First N elements in self that are already sorted and not duplicated. + let mut sorted_len = self.len(); + // Use SmallVec::extend as it has a specialized code for slice and trusted-len iterators. + self.extend(other); + // `sorted_len` is the length of the sorted run before extension + // window slice `w` is guaranteed to have a length of 2. + #[allow(clippy::indexing_slicing)] + { + // Count new elements that are sorted and non-duplicated. + // Starting from the end of the existing sorted run, if any. + // Thus, start the slice at sorted_len.saturating_sub(1). + sorted_len += self[sorted_len.saturating_sub(1)..] + .windows(2) + .take_while(|w| w[0].0 < w[1].0) + .count(); + } + // `windows(2)` only yields `slice len - 1` times, or none if the slice is empty. + // In other words, the first extended element of the slice won't be counted as sorted + // if self was initially empty (sorted_len == 0). We adjust this by adding 1 if the + // original slice was empty but became not empty after extend. + sorted_len += (sorted_len == 0 && !self.is_empty()) as usize; + + // If everything was in order, we're done + if sorted_len >= self.len() { + return; + } + + // Use stable sort to keep relative order of duplicates. + self.sort_by(|a, b| a.0.cmp(&b.0)); + // Deduplicate by keeping the last element of the run in the first slice. + let (dedup, _merged_dup) = partition_dedup_by(self); + sorted_len = dedup.len(); + self.truncate(sorted_len); + } +} + +/// This is copied from `litemap` crate, since this is not exported. +/// See +/// +/// Moves all but the _last_ of consecutive elements to the end of the slice satisfying +/// equality on K. +/// +/// Returns two slices. The first contains no consecutive repeated elements. +/// The second contains all the duplicates in no specified order. +/// +/// This is based on std::slice::partition_dedup_by (currently unstable) but retains the +/// _last_ element of the duplicate run in the first slice (instead of first). +#[cfg(feature = "litemap")] +#[inline] +#[expect(clippy::type_complexity)] +fn partition_dedup_by(v: &mut [(K, V)]) -> (&mut [(K, V)], &mut [(K, V)]) { + if v.len() <= 1 { + return (v, &mut []); + } + + let mut read_idx: usize = 1; + let mut write_idx: usize = 1; + + while let Some((before_read, [read, ..])) = v.split_at_mut_checked(read_idx) { + // First, `read_idx >= write_idx` is always true as `read_idx` is always incremented + // whereas `write_idx` is only incremented when a distinct element is found. + // Second, before_read is always at least 1 length due to read_idx being initialized to 1. + // Thus it is safe to index before_read with `write_idx - 1`. + #[expect(clippy::indexing_slicing)] + let prev_write = &mut before_read[write_idx - 1]; + if read.0 == prev_write.0 { + core::mem::swap(read, prev_write); + } else { + // Equivalent to checking if write_idx == read_idx + if let Some(write) = before_read.get_mut(write_idx) { + core::mem::swap(read, write); + } + write_idx += 1; + } + read_idx += 1; + } + v.split_at_mut(write_idx) +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreConstEmpty for SmallVec<(K, V), N> { + const EMPTY: Self = SmallVec::new(); +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreFromIterable for SmallVec<(K, V), N> { + #[inline] + fn lm_sort_from_iter>(iter: I) -> Self { + let mut v = Self::new(); + v.lm_extend(iter); + v + } +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreFromIterator for SmallVec<(K, V), N> {} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl StoreIntoIterator for SmallVec<(K, V), N> { + type KeyValueIntoIter = IntoIter<(K, V), N>; + + #[inline] + fn lm_into_iter(self) -> Self::KeyValueIntoIter { + self.into_iter() + } + + #[inline] + fn lm_extend_end(&mut self, other: Self) { + self.extend(other); + } + + #[inline] + fn lm_extend_start(&mut self, other: Self) { + self.splice(0..0, other); + } +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl<'a, K: 'a, V: 'a, const N: usize> StoreIterable<'a, K, V> for SmallVec<(K, V), N> { + type KeyValueIter = core::iter::Map, MapF>; + + #[inline] + fn lm_iter(&'a self) -> Self::KeyValueIter { + self.iter().map(map_f) + } +} + +#[cfg(feature = "litemap")] +#[cfg_attr(docsrs, doc(cfg(feature = "litemap")))] +impl<'a, K: 'a, V: 'a, const N: usize> StoreIterableMut<'a, K, V> for SmallVec<(K, V), N> { + type KeyValueIterMut = core::iter::Map, MapFMut>; + + #[inline] + fn lm_iter_mut(&'a mut self) -> Self::KeyValueIterMut { + self.iter_mut().map(map_f_mut) + } +} diff --git a/src/tests.rs b/src/tests.rs index 2082f96..ed70736 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1121,3 +1121,10 @@ mod buf_mut { } } } + +// Note: To test against the `litemap` feature, you need to enable `litemap/testing` feature in cli. +#[cfg(all(test, feature = "litemap"))] +#[test] +fn test_litemap() { + litemap::testing::check_store_full::>() +}