TryPush trait to enable non-reallocating pushes

Signed-off-by: Moritz Hoffmann <[email protected]>

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<S, L> OffsetContainer<usize> for OffsetList<S, L>
@@ -407,6 +420,16 @@ where
     fn heap_size<F: FnMut(usize, usize)>(&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<S, L> OffsetContainer<usize> for OffsetOptimized<S, L>
@@ -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);
+            }
         }
     }
 

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<T, S, const N: usize> TryPush<[T; N]> for OwnedRegion<T, S>
+where
+    [T]: ToOwned,
+    S: Storage<T>
+        + for<'a> PushStorage<PushIter<[T; N]>>
+        + std::ops::Index<std::ops::Range<usize>, Output = [T]>,
+{
+    #[inline]
+    fn try_push(&mut self, item: [T; N]) -> Result<<OwnedRegion<T> 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<T, S, const N: usize> Push<&[T; N]> for OwnedRegion<T, S>
 where
     T: Clone,
@@ -141,6 +165,27 @@ where
     }
 }
 
+impl<T, S, const N: usize> TryPush<&[T; N]> for OwnedRegion<T, S>
+where
+    T: Clone,
+    S: Storage<T>
+        + for<'a> PushStorage<&'a [T]>
+        + std::ops::Index<std::ops::Range<usize>, Output = [T]>,
+{
+    #[inline]
+    fn try_push<'a>(
+        &mut self,
+        item: &'a [T; N],
+    ) -> Result<<OwnedRegion<T> 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<T, S, const N: usize> Push<&&[T; N]> for OwnedRegion<T, S>
 where
     T: Clone,
@@ -183,6 +228,33 @@ where
     }
 }
 
+impl<T, S> TryPush<&[T]> for OwnedRegion<T, S>
+where
+    T: Clone,
+    S: Storage<T>
+        + for<'a> PushStorage<&'a [T]>
+        + std::ops::Index<std::ops::Range<usize>, Output = [T]>,
+{
+    #[inline]
+    fn try_push<'a>(
+        &mut self,
+        item: &'a [T],
+    ) -> Result<<OwnedRegion<T, S> 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<T: Clone, S: Storage<T>> Push<&&[T]> for OwnedRegion<T, S>
 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 = <OwnedRegion<u8>>::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()));
+    }
 }

src/impls/storage.rs

@@ -48,6 +48,10 @@ pub trait Storage<T>: 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<T> Storage<T> for Vec<T> {
@@ -83,6 +87,11 @@ impl<T> Storage<T> for Vec<T> {
     fn is_empty(&self) -> bool {
         self.is_empty()
     }
+
+    #[inline]
+    fn capacity(&self) -> usize {
+        self.capacity()
+    }
 }
 
 /// Push an item into storage.

src/lib.rs

@@ -112,6 +112,17 @@ pub trait Push<T>: Region {
     fn push(&mut self, item: T) -> Self::Index;
 }
 
+/// Push an item `T` into a region.
+pub trait TryPush<T>: 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<Self::Index, T>;
+
+    /// 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<R: Region, S: OffsetContainer<<R as Region>::Index>> FlatStack<R, S> {
         self.indices.push(index);
     }
 
+    /// Appends the element to the back of the stack, if there is sufficient capacity
+    #[inline]
+    pub fn try_push<T>(&mut self, item: T) -> Result<(), T>
+    where
+        R: TryPush<T>,
+    {
+        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<T>(&mut self, item: &T) -> bool
+    where
+        R: TryPush<T>,
+    {
+        // TODO: Include `indices` in the check.
+        self.region.can_push(item)
+    }
+
     /// Returns the element at the `offset` position.
     #[inline]
     #[must_use]