From 23874294ae01116f591be183bbc6babfdb63467c Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Tue, 2 Jul 2024 22:13:25 -0400 Subject: [PATCH] TryPush trait to enable non-reallocating pushes Signed-off-by: Moritz Hoffmann --- src/impls/offsets.rs | 46 +++++++++++++++++++-- src/impls/slice_owned.rs | 86 +++++++++++++++++++++++++++++++++++++++- src/impls/storage.rs | 9 +++++ src/lib.rs | 32 +++++++++++++++ 4 files changed, 169 insertions(+), 4 deletions(-) diff --git a/src/impls/offsets.rs b/src/impls/offsets.rs index c968c24..3bf72a9 100644 --- a/src/impls/offsets.rs +++ b/src/impls/offsets.rs @@ -248,7 +248,11 @@ where /// Reserve space for `additional` elements. #[inline] pub fn reserve(&mut self, additional: usize) { - self.smol.reserve(additional); + if self.chonk.is_empty() { + self.smol.reserve(additional); + } else { + self.chonk.reserve(additional); + } } /// Remove all elements. @@ -299,6 +303,15 @@ where fn is_empty(&self) -> bool { self.is_empty() } + + #[inline] + fn capacity(&self) -> usize { + if self.chonk.is_empty() { + self.smol.capacity() + } else { + self.chonk.capacity() + } + } } impl OffsetContainer for OffsetList @@ -407,6 +420,16 @@ where fn heap_size(&self, callback: F) { self.spilled.heap_size(callback); } + + #[inline] + fn capacity(&self) -> usize { + if self.spilled.is_empty() { + // TODO: What else could we do here? `usize::MAX`? + 1 + } else { + self.spilled.capacity() + } + } } impl OffsetContainer for OffsetOptimized @@ -439,8 +462,25 @@ where where I::IntoIter: ExactSizeIterator, { - for item in iter { - self.push(item); + let mut iter = iter.into_iter(); + if !self.spilled.is_empty() { + // We certainly push into `spilled`, so reserve enough space. + let (lower, upper) = iter.size_hint(); + self.spilled.reserve(upper.unwrap_or(lower)); + } + while let Some(item) = iter.next() { + if self.spilled.is_empty() { + let inserted = self.strided.push(item); + if !inserted { + // This is a transition point from pushing into `strided` to pushing into + // `spilled`. + let (lower, upper) = iter.size_hint(); + self.spilled.reserve(upper.unwrap_or(lower)); + self.spilled.push(item); + } + } else { + self.spilled.push(item); + } } } diff --git a/src/impls/slice_owned.rs b/src/impls/slice_owned.rs index 209e0e0..2f34d9f 100644 --- a/src/impls/slice_owned.rs +++ b/src/impls/slice_owned.rs @@ -6,7 +6,7 @@ use std::marker::PhantomData; use serde::{Deserialize, Serialize}; use crate::impls::storage::{PushStorage, Storage}; -use crate::{Push, PushIter, Region, ReserveItems}; +use crate::{Push, PushIter, Region, ReserveItems, TryPush}; /// A container for owned types. /// @@ -128,6 +128,30 @@ where } } +impl TryPush<[T; N]> for OwnedRegion +where + [T]: ToOwned, + S: Storage + + for<'a> PushStorage> + + std::ops::Index, Output = [T]>, +{ + #[inline] + fn try_push(&mut self, item: [T; N]) -> Result< as Region>::Index, [T; N]> { + if self.can_push(&item) { + let start = self.slices.len(); + self.slices.push_storage(PushIter(item)); + Ok((start, self.slices.len())) + } else { + Err(item) + } + } + + #[inline] + fn can_push(&self, item: &[T; N]) -> bool { + self.slices.capacity() - self.slices.len() >= item.len() + } +} + impl Push<&[T; N]> for OwnedRegion where T: Clone, @@ -141,6 +165,27 @@ where } } +impl TryPush<&[T; N]> for OwnedRegion +where + T: Clone, + S: Storage + + for<'a> PushStorage<&'a [T]> + + std::ops::Index, Output = [T]>, +{ + #[inline] + fn try_push<'a>( + &mut self, + item: &'a [T; N], + ) -> Result< as Region>::Index, &'a [T; N]> { + self.try_push(item.as_slice()).map_err(|_| item) + } + + #[inline] + fn can_push(&self, item: &&[T; N]) -> bool { + self.can_push(&item.as_slice()) + } +} + impl Push<&&[T; N]> for OwnedRegion where T: Clone, @@ -183,6 +228,33 @@ where } } +impl TryPush<&[T]> for OwnedRegion +where + T: Clone, + S: Storage + + for<'a> PushStorage<&'a [T]> + + std::ops::Index, Output = [T]>, +{ + #[inline] + fn try_push<'a>( + &mut self, + item: &'a [T], + ) -> Result< as Region>::Index, &'a [T]> { + if self.can_push(&item) { + let start = self.slices.len(); + self.slices.push_storage(item); + Ok((start, self.slices.len())) + } else { + Err(item) + } + } + + #[inline] + fn can_push(&self, item: &&[T]) -> bool { + self.slices.capacity() - self.slices.len() >= item.len() + } +} + impl> Push<&&[T]> for OwnedRegion where for<'a> Self: Push<&'a [T]>, @@ -322,4 +394,16 @@ mod tests { let index = r.push(PushIter(iter)); assert_eq!([1, 1, 1, 1], r.index(index)); } + + #[test] + fn try_push() { + let mut r = >::default(); + assert!(!r.can_push(&[1; 4])); + assert_eq!(r.try_push(&[1; 4]), Err(&[1; 4])); + r.reserve_items(std::iter::once(&[1; 4])); + assert!(r.can_push(&[1; 4])); + let index = r.try_push(&[1; 4]); + assert_eq!(index, Ok((0, 4))); + assert_eq!([1, 1, 1, 1], r.index(index.unwrap())); + } } diff --git a/src/impls/storage.rs b/src/impls/storage.rs index 6e1b208..1831beb 100644 --- a/src/impls/storage.rs +++ b/src/impls/storage.rs @@ -48,6 +48,10 @@ pub trait Storage: Default { /// Returns `true` if empty, i.e., it doesn't contain any elements. #[must_use] fn is_empty(&self) -> bool; + + /// Returns the capacity of the storage. + #[must_use] + fn capacity(&self) -> usize; } impl Storage for Vec { @@ -83,6 +87,11 @@ impl Storage for Vec { fn is_empty(&self) -> bool { self.is_empty() } + + #[inline] + fn capacity(&self) -> usize { + self.capacity() + } } /// Push an item into storage. diff --git a/src/lib.rs b/src/lib.rs index 23af4ae..d892827 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -112,6 +112,17 @@ pub trait Push: Region { fn push(&mut self, item: T) -> Self::Index; } +/// Push an item `T` into a region. +pub trait TryPush: Region { + /// Push `item` into self, returning an index that allows to look up the + /// corresponding read item. + fn try_push(&mut self, item: T) -> Result; + + /// Test if an item can be pushed into the region without reallocation. + #[must_use] + fn can_push(&self, item: &T) -> bool; +} + /// Reserve space in the receiving region. /// /// Closely related to [`Push`], but separate because target type is likely different. @@ -226,6 +237,27 @@ impl::Index>> FlatStack { self.indices.push(index); } + /// Appends the element to the back of the stack, if there is sufficient capacity + #[inline] + pub fn try_push(&mut self, item: T) -> Result<(), T> + where + R: TryPush, + { + let index = self.region.try_push(item)?; + self.indices.push(index); + Ok(()) + } + + /// Appends the element to the back of the stack, if there is sufficient capacity + #[inline] + pub fn can_push(&mut self, item: &T) -> bool + where + R: TryPush, + { + // TODO: Include `indices` in the check. + self.region.can_push(item) + } + /// Returns the element at the `offset` position. #[inline] #[must_use]