Add IntoIter implementations for ArrayQueue & SegQueue

crossbeam-queue/src/array_queue.rs

@@ -8,6 +8,7 @@ use core::cell::UnsafeCell;
 use core::fmt;
 use core::marker::PhantomData;
 use core::mem::MaybeUninit;
+use core::ptr;
 use core::sync::atomic::{self, AtomicUsize, Ordering};
 
 use crossbeam_utils::{Backoff, CachePadded};
@@ -430,3 +431,52 @@ impl<T> fmt::Debug for ArrayQueue<T> {
         f.pad("ArrayQueue { .. }")
     }
 }
+
+impl<T> IntoIterator for ArrayQueue<T> {
+    type Item = T;
+
+    type IntoIter = IntoIter<T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        IntoIter { value: self }
+    }
+}
+
+#[derive(Debug)]
+pub struct IntoIter<T> {
+    value: ArrayQueue<T>,
+}
+
+impl<T> Iterator for IntoIter<T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let value = &mut self.value;
+        let head = *value.head.get_mut();
+        if value.head.get_mut() != value.tail.get_mut() {
+            let index = head & (value.one_lap - 1);
+            let lap = head & !(value.one_lap - 1);
+            // SAFETY: We have mutable access to this, so we can read without
+            // worrying about concurrency. Furthermore, we know this is
+            // initialized because it is the value pointed at by `value.head`
+            // and this is a non-empty queue.
+            let val = unsafe {
+                let slot = &mut *value.buffer.add(index);
+                ptr::read(slot.value.get()).assume_init()
+            };
+            let new = if index + 1 < value.cap {
+                // Same lap, incremented index.
+                // Set to `{ lap: lap, index: index + 1 }`.
+                head + 1
+            } else {
+                // One lap forward, index wraps around to zero.
+                // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
+                lap.wrapping_add(value.one_lap)
+            };
+            *value.head.get_mut() = new;
+            Option::Some(val)
+        } else {
+            Option::None
+        }
+    }
+}

crossbeam-queue/src/seg_queue.rs

@@ -484,3 +484,62 @@ impl<T> Default for SegQueue<T> {
         SegQueue::new()
     }
 }
+
+impl<T> IntoIterator for SegQueue<T> {
+    type Item = T;
+
+    type IntoIter = IntoIter<T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        IntoIter { value: self }
+    }
+}
+
+#[derive(Debug)]
+pub struct IntoIter<T> {
+    value: SegQueue<T>,
+}
+
+impl<T> Iterator for IntoIter<T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let value = &mut self.value;
+        let head = *value.head.index.get_mut();
+        let tail = *value.tail.index.get_mut();
+        if head >> SHIFT == tail >> SHIFT {
+            None
+        } else {
+            let block = *value.head.block.get_mut();
+            let offset = (head >> SHIFT) % LAP;
+
+            // SAFETY: We have mutable access to this, so we can read without
+            // worrying about concurrency. Furthermore, we know this is
+            // initialized because it is the value pointed at by `value.head`
+            // and this is a non-empty queue.
+            let item = unsafe {
+                let slot = (*block).slots.get_unchecked(offset);
+                let p = &mut *slot.value.get();
+                p.as_mut_ptr().read()
+            };
+            if offset + 1 == BLOCK_CAP {
+                // Deallocate the block and move to the next one.
+                // SAFETY: The block is initialized because we've been reading
+                // from it this entire time. We can drop it b/c everything has
+                // been read out of it, so nothing is pointing to it anymore.
+                unsafe {
+                    let next = *(*block).next.get_mut();
+                    drop(Box::from_raw(block));
+                    *value.head.block.get_mut() = next;
+                }
+                // The last value in a block is empty, so skip it
+                *value.head.index.get_mut() = head.wrapping_add(2 << SHIFT);
+                // Double-check that we're pointing to the first item in a block.
+                debug_assert_eq!((*value.head.index.get_mut() >> SHIFT) % LAP, 0);
+            } else {
+                *value.head.index.get_mut() = head.wrapping_add(1 << SHIFT);
+            }
+            Some(item)
+        }
+    }
+}

crossbeam-queue/tests/array_queue.rs

@@ -255,3 +255,14 @@ fn linearizable() {
     })
     .unwrap();
 }
+
+#[test]
+fn into_iter() {
+    let q = ArrayQueue::new(100);
+    for i in 0..100 {
+        q.push(i).unwrap();
+    }
+    for (i, j) in q.into_iter().enumerate() {
+        assert_eq!(i, j);
+    }
+}

crossbeam-queue/tests/seg_queue.rs

@@ -166,3 +166,25 @@ fn drops() {
         assert_eq!(DROPS.load(Ordering::SeqCst), steps + additional);
     }
 }
+
+#[test]
+fn into_iter() {
+    let q = SegQueue::new();
+    for i in 0..100 {
+        q.push(i);
+    }
+    for (i, j) in q.into_iter().enumerate() {
+        assert_eq!(i, j);
+    }
+}
+
+#[test]
+fn into_iter_drop() {
+    let q = SegQueue::new();
+    for i in 0..100 {
+        q.push(i);
+    }
+    for (i, j) in q.into_iter().enumerate().take(50) {
+        assert_eq!(i, j);
+    }
+}