From 93ce5b2a590a67b548f8a28ea813b2a6841916c6 Mon Sep 17 00:00:00 2001
From: Lucas Van Laer <lucas.van.laer@gmail.com>
Date: Mon, 21 Jul 2025 13:25:55 +0200
Subject: [PATCH] Adds `as_slice` and friends to basic slice and vec iterators

This commit adds the methods `as_slice`, `as_mut_slice`, `into_slice`
and `AsRef<[T]>` on the appropriate iterators for slice iterators and
vec iterators.
This allows custom producers to use these iterators to decide on
custom split points based on the remaining data in the slice.
---
 src/slice/mod.rs | 77 ++++++++++++++++++++++++++++++++++++++++++++++++
 src/vec.rs       | 74 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 151 insertions(+)

diff --git a/src/slice/mod.rs b/src/slice/mod.rs
index 25e911663..bf624d8a1 100644
--- a/src/slice/mod.rs
+++ b/src/slice/mod.rs
@@ -785,6 +785,28 @@ pub struct Iter<'data, T: Sync> {
     slice: &'data [T],
 }
 
+impl<'data, T: Sync> Iter<'data, T> {
+    /// Views the underlying data as a subslice of the original data.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let slice = &['a', 'b', 'c'];
+    /// let iter = slice[1..3].par_iter();
+    /// assert_eq!(iter.as_slice(), &['b', 'c']);
+    /// ```
+    pub fn as_slice(&self) -> &'data [T] {
+        self.slice
+    }
+}
+
+impl<'data, T: Sync> AsRef<[T]> for Iter<'data, T> {
+    fn as_ref(&self) -> &[T] {
+        self.as_slice()
+    }
+}
+
 impl<'data, T: Sync> Clone for Iter<'data, T> {
     fn clone(&self) -> Self {
         Iter { ..*self }
@@ -932,6 +954,61 @@ pub struct IterMut<'data, T: Send> {
     slice: &'data mut [T],
 }
 
+impl<'data, T: Send> IterMut<'data, T> {
+    /// Returns the remaining items in this iterator as a mutable slice with the original lifetime.
+    ///
+    /// To return with the lifetime of the original slice, the iterator must be consumed.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut slice = &mut [1, 2, 3];
+    /// let mut iter = slice.par_iter_mut();
+    /// let mut slice_mut = iter.into_slice();
+    /// slice_mut[0] = 7;
+    /// assert_eq!(slice_mut, &[7, 2, 3]);
+    /// ```
+    pub fn into_slice(self) -> &'data mut [T] {
+        self.slice
+    }
+
+    /// Views the underlying data as a subslice of the original data.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut slice = &[1, 2, 3];
+    /// let mut iter = slice[1..].par_iter();
+    /// assert_eq!(iter.as_slice(), &[2, 3]);
+    /// ```
+    pub fn as_slice(&self) -> &[T] {
+        self.slice
+    }
+
+    /// Views the underlying data as a mutable subslice of the original data.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut slice = &mut [1, 2, 3];
+    /// let mut iter = slice[1..].par_iter_mut();
+    /// iter.as_mut_slice()[0] = 7;
+    /// assert_eq!(iter.as_mut_slice(), &[7, 3]);
+    /// ```
+    pub fn as_mut_slice(&mut self) -> &mut [T] {
+        self.slice
+    }
+}
+
+impl<'data, T: Send> AsRef<[T]> for IterMut<'data, T> {
+    fn as_ref(&self) -> &[T] {
+        self.as_slice()
+    }
+}
+
 impl<'data, T: Send + 'data> ParallelIterator for IterMut<'data, T> {
     type Item = &'data mut T;
 
diff --git a/src/vec.rs b/src/vec.rs
index 7892f53fa..36d2ad395 100644
--- a/src/vec.rs
+++ b/src/vec.rs
@@ -39,6 +39,43 @@ pub struct IntoIter<T: Send> {
     vec: Vec<T>,
 }
 
+impl<T: Send> IntoIter<T> {
+    /// Returns the remaining items of this iterator as a slice.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut vec = vec!['a', 'b', 'c'];
+    /// let drain = vec.par_drain(0..2);
+    /// assert_eq!(drain.as_slice(), &['a', 'b']);
+    /// ```
+    pub fn as_slice(&self) -> &[T] {
+        &self.vec
+    }
+
+    /// Returns the remaining items of this iterator as a mutable slice.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut vec = vec!['a', 'b', 'c'];
+    /// let mut drain = vec.par_drain(0..2);
+    /// *drain.as_mut_slice().last_mut().unwrap() = 'd';
+    /// assert_eq!(drain.as_slice(), &['a', 'd']);
+    /// ```
+    pub fn as_mut_slice(&mut self) -> &mut [T] {
+        &mut self.vec
+    }
+}
+
+impl<T: Send> AsRef<[T]> for IntoIter<T> {
+    fn as_ref(&self) -> &[T] {
+        self.as_slice()
+    }
+}
+
 impl<T: Send> IntoParallelIterator for Vec<T> {
     type Item = T;
     type Iter = IntoIter<T>;
@@ -105,6 +142,43 @@ pub struct Drain<'data, T: Send> {
     orig_len: usize,
 }
 
+impl<'data, T: Send> Drain<'data, T> {
+    /// Returns the remaining items of this iterator as a slice.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut vec = vec!['a', 'b', 'c'];
+    /// let drain = vec.par_drain(0..2);
+    /// assert_eq!(drain.as_slice(), &['a', 'b']);
+    /// ```
+    pub fn as_slice(&self) -> &[T] {
+        &self.vec[self.range.clone()]
+    }
+
+    /// Returns the remaining items of this iterator as a mutable slice.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use rayon::prelude::*;
+    /// let mut vec = vec!['a', 'b', 'c'];
+    /// let mut drain = vec.par_drain(0..2);
+    /// *drain.as_mut_slice().last_mut().unwrap() = 'd';
+    /// assert_eq!(drain.as_slice(), &['a', 'd']);
+    /// ```
+    pub fn as_mut_slice(&mut self) -> &mut [T] {
+        &mut self.vec[self.range.clone()]
+    }
+}
+
+impl<'data, T: Send> AsRef<[T]> for Drain<'data, T> {
+    fn as_ref(&self) -> &[T] {
+        self.as_slice()
+    }
+}
+
 impl<'data, T: Send> ParallelIterator for Drain<'data, T> {
     type Item = T;