queue: add drain for SegQueue

crossbeam-queue/src/seg_queue.rs

@@ -4,7 +4,8 @@ use core::{
     cell::UnsafeCell,
     fmt,
     marker::PhantomData,
-    mem::MaybeUninit,
+    mem::{self, MaybeUninit},
+    ops::{Bound, RangeBounds},
     panic::{RefUnwindSafe, UnwindSafe},
     ptr,
     sync::atomic::{self, AtomicPtr, AtomicUsize, Ordering},
@@ -525,6 +526,99 @@ impl<T> SegQueue<T> {
         }
     }
 
+    /// Returns an iterator that drains elements from the queue within the given range.
+    ///
+    /// Elements before the range and after the range remain in the queue.
+    /// If the `Drain` iterator is dropped before being fully consumed,
+    /// the remaining elements within the range are dropped, and elements
+    /// outside the range are preserved in their original order.
+    ///
+    /// If `mem::forget` is called on the `Drain`, the queue is left empty.
+    /// All elements not yet yielded are leaked, but the queue remains in a
+    /// consistent (empty) state — no corruption occurs.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the start of the range is greater than the end.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use crossbeam_queue::SegQueue;
+    ///
+    /// // Full drain
+    /// let mut q = SegQueue::new();
+    /// for i in 0..5 { q.push(i); }
+    /// let v: Vec<_> = q.drain(..).collect();
+    /// assert_eq!(v, [0, 1, 2, 3, 4]);
+    /// assert!(q.is_empty());
+    ///
+    /// // Prefix drain
+    /// let mut q = SegQueue::new();
+    /// for i in 0..5 { q.push(i); }
+    /// let v: Vec<_> = q.drain(..3).collect();
+    /// assert_eq!(v, [0, 1, 2]);
+    /// assert_eq!(q.len(), 2);
+    ///
+    /// // Range drain
+    /// let mut q = SegQueue::new();
+    /// for i in 0..5 { q.push(i); }
+    /// let v: Vec<_> = q.drain(1..4).collect();
+    /// assert_eq!(v, [1, 2, 3]);
+    /// assert_eq!(q.len(), 2);
+    /// ```
+    pub fn drain<R: RangeBounds<usize>>(&mut self, range: R) -> Drain<'_, T> {
+        let start = match range.start_bound() {
+            Bound::Included(&n) => n,
+            Bound::Excluded(&n) => n.checked_add(1).expect("start index overflow"),
+            Bound::Unbounded => 0,
+        };
+        let end = match range.end_bound() {
+            Bound::Included(&n) => n.checked_add(1).expect("end index overflow"),
+            Bound::Excluded(&n) => n,
+            Bound::Unbounded => usize::MAX,
+        };
+        assert!(start <= end, "drain range start is greater than end");
+
+        // Step 1: pop `start` elements into prefix via pop_mut.
+        // prefix owns its own blocks — no pointer sharing with rest.
+        let mut prefix = SegQueue::new();
+        let mut actual_start = 0;
+        for _ in 0..start {
+            match self.pop_mut() {
+                Some(v) => {
+                    prefix.push_mut(v);
+                    actual_start += 1;
+                }
+                None => break,
+            }
+        }
+
+        // Step 2: snapshot the remaining queue (drain range + suffix) into `rest`
+        // by transferring head/tail pointers. No element copying
+        let mut rest = SegQueue::new();
+        *rest.head.index.get_mut() = *self.head.index.get_mut();
+        *rest.head.block.get_mut() = *self.head.block.get_mut();
+        *rest.tail.index.get_mut() = *self.tail.index.get_mut();
+        *rest.tail.block.get_mut() = *self.tail.block.get_mut();
+
+        // Step 3: reset original queue to empty immediately.
+        // If mem::forget is called on Drain, queue is left empty — consistent but leaky.
+        *self.head.index.get_mut() = 0;
+        *self.head.block.get_mut() = ptr::null_mut();
+        *self.tail.index.get_mut() = 0;
+        *self.tail.block.get_mut() = ptr::null_mut();
+
+        let remaining = end.saturating_sub(actual_start);
+
+        Drain {
+            queue: self,
+            prefix,
+            rest,
+            remaining,
+        }
+    }
+
     /// Returns `true` if the queue is empty.
     ///
     /// # Examples
@@ -596,6 +690,72 @@ impl<T> SegQueue<T> {
     }
 }
 
+/// A draining iterator for `SegQueue<T>`.
+///
+/// This struct is created by the [`drain`] method on [`SegQueue`].
+///
+/// [`drain`]: SegQueue::drain
+pub struct Drain<'a, T> {
+    /// The original queue, reset to empty at `Drain` creation.
+    /// Restored to prefix + suffix on `Drop`.
+    queue: &'a mut SegQueue<T>,
+    /// Saved prefix elements (before the drain range), owns its own blocks.
+    prefix: SegQueue<T>,
+    /// Drain range + suffix, owned via pointer snapshot from original queue.
+    rest: SegQueue<T>,
+    /// How many elements are left to yield/drop within the drain range.
+    remaining: usize,
+}
+
+impl<T> Iterator for Drain<'_, T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        if self.remaining == 0 {
+            return None;
+        }
+        let val = self.rest.pop_mut();
+        if val.is_some() {
+            self.remaining -= 1;
+        }
+        val
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // Credit: size_hint and ExactSizeIterator suggested by @laycookie
+        // in https://github.com/crossbeam-rs/crossbeam/issues/1228
+        let remaining = self.remaining.min(self.rest.len());
+        (remaining, Some(remaining))
+    }
+}
+
+impl<T> ExactSizeIterator for Drain<'_, T> {}
+
+impl<T> Drop for Drain<'_, T> {
+    fn drop(&mut self) {
+        // Step 1: drop remaining elements in the drain range.
+        while self.remaining > 0 {
+            if self.next().is_none() {
+                break;
+            }
+        }
+        // Step 2: move suffix (whatever remains in rest) into prefix.
+        // At this point rest contains only suffix elements.
+        while let Some(v) = self.rest.pop_mut() {
+            self.prefix.push_mut(v);
+        }
+        // Step 3: swap prefix (now prefix + suffix) back into original queue.
+        mem::swap(self.queue, &mut self.prefix);
+        // self.prefix (now the old empty queue) drops cleanly here.
+    }
+}
+
+impl<T> fmt::Debug for Drain<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.pad("Drain { .. }")
+    }
+}
+
 impl<T> Drop for SegQueue<T> {
     fn drop(&mut self) {
         let mut head = *self.head.index.get_mut();