feat: added BumpBox<[T]>::partition

.vscode/cspell.json

@@ -32,6 +32,7 @@
     "emscripten",
     "expl",
     "foobarbaz",
+    "genericity",
     "haha",
     "hashbrown",
     "Hasher",
@@ -61,6 +62,7 @@
     "rangemap",
     "rdme",
     "repr",
+    "rfind",
     "rustc",
     "rustdoc",
     "rustfmt",

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 - **added**: `from_utf8_unchecked` for all string types
 - **added**: `BumpString::into_fixed_string`
+- **added**: `BumpBox<[T]>::partition`
 
 ## 0.5.6 (2024-08-17)
 - **fixed:** `alloc_iter` and `alloc_fmt` to not take up more than the necessary memory

src/bump_box.rs

@@ -1535,7 +1535,7 @@ impl<'a, T> BumpBox<'a, [T]> {
         }
     }
 
-    /// Consumes an iterator, creating two collections from it.
+    /// Consumes `self`, creating two boxed slices from it.
     ///
     /// The predicate passed to `partition()` can return `true`, or `false`.
     /// `partition()` returns a pair, all of the elements for which it returned
@@ -1553,19 +1553,18 @@ impl<'a, T> BumpBox<'a, [T]> {
     /// ```
     /// # use bump_scope::Bump;
     /// # let bump: Bump = Bump::new();
-    /// let slice = bump.alloc_slice_copy(&[1, 2, 3]);
+    /// let slice = bump.alloc_slice_copy(&[1, 2, 3, 4, 5, 6, 7]);
     ///
     /// let (even, odd) = slice.partition(|n| n % 2 == 0);
     ///
-    /// assert_eq!(even, vec![2]);
-    /// assert_eq!(odd, vec![1, 3]);
+    /// assert!(even.iter().all(|n| n % 2 == 0));
+    /// assert!(odd.iter().all(|n| n % 2 != 0));
     /// ```
-    #[cfg(test)]
-    pub fn partition<F>(mut self, f: F) -> (BumpBox<'a, [T]>, BumpBox<'a, [T]>)
+    pub fn partition<F>(mut self, f: F) -> (Self, Self)
     where
         F: FnMut(&T) -> bool,
     {
-        let index = self.iter_mut().partition_in_place(f);
+        let index = polyfill::iter::partition_in_place(self.iter_mut(), f);
         self.split_at(index)
     }
 }

src/polyfill/iter.rs

@@ -0,0 +1,41 @@
+pub(crate) fn partition_in_place<'a, T: 'a, P>(
+    mut iter: impl DoubleEndedIterator<Item = &'a mut T>,
+    ref mut predicate: P,
+) -> usize
+where
+    P: FnMut(&T) -> bool,
+{
+    // FIXME: should we worry about the count overflowing? The only way to have more than
+    // `usize::MAX` mutable references is with ZSTs, which aren't useful to partition...
+
+    // These closure "factory" functions exist to avoid genericity in `Self`.
+
+    #[inline]
+    fn is_false<'a, T>(
+        predicate: &'a mut impl FnMut(&T) -> bool,
+        true_count: &'a mut usize,
+    ) -> impl FnMut(&&mut T) -> bool + 'a {
+        move |x| {
+            let p = predicate(&**x);
+            *true_count += p as usize;
+            !p
+        }
+    }
+
+    #[inline]
+    fn is_true<T>(predicate: &mut impl FnMut(&T) -> bool) -> impl FnMut(&&mut T) -> bool + '_ {
+        move |x| predicate(&**x)
+    }
+
+    // Repeatedly find the first `false` and swap it with the last `true`.
+    let mut true_count = 0;
+    while let Some(head) = iter.find(is_false(predicate, &mut true_count)) {
+        if let Some(tail) = iter.rfind(is_true(predicate)) {
+            crate::mem::swap(head, tail);
+            true_count += 1;
+        } else {
+            break;
+        }
+    }
+    true_count
+}

src/polyfill/mod.rs

@@ -1,6 +1,7 @@
 //! A collection of stuff that is missing from stable std.
 //! Most of this is from nightly.
 
+pub mod iter;
 pub mod nonnull;
 pub mod nonzero;
 pub mod pointer;