A few optimizations

Author: mgeier

src/chunks.rs

@@ -349,14 +349,13 @@ where
 {
     /// Fills all slots with the [`Default`] value.
     fn from(chunk: WriteChunkUninit<'a, T>) -> Self {
-        for slot in chunk
-            .first_slice
-            .iter_mut()
-            .chain(chunk.second_slice.iter_mut())
-        {
-            unsafe {
-                slot.as_mut_ptr().write(Default::default());
-            }
+        // NB: Using Iterator::chain() to iterate over both slices
+        //     led to worse optimization (with rustc 1.57).
+        for slot in chunk.first_slice.iter_mut() {
+            *slot = MaybeUninit::new(Default::default());
+        }
+        for slot in chunk.second_slice.iter_mut() {
+            *slot = MaybeUninit::new(Default::default());
         }
         WriteChunk(Some(chunk))
     }
@@ -558,18 +557,33 @@ impl<T> WriteChunkUninit<'_, T> {
     where
         I: IntoIterator<Item = T>,
     {
-        let iterated = self
-            .first_slice
-            .iter_mut()
-            .chain(self.second_slice.iter_mut())
-            .zip(iter)
-            .map(|(slot, item)| {
-                // Safety: It is allowed to write to this memory slot
-                unsafe {
-                    slot.as_mut_ptr().write(item);
+        let mut iter = iter.into_iter();
+        let mut iterated = 0;
+        // NB: Iterating over slices (instead of using pointers)
+        //     led to worse optimization (with rustc 1.57).
+        'outer: for &(ptr, len) in &[
+            (
+                self.first_slice.as_mut_ptr().cast::<T>(),
+                self.first_slice.len(),
+            ),
+            (
+                self.second_slice.as_mut_ptr().cast::<T>(),
+                self.second_slice.len(),
+            ),
+        ] {
+            for i in 0..len {
+                match iter.next() {
+                    Some(item) => {
+                        // Safety: It is allowed to write to this memory slot
+                        unsafe {
+                            ptr.add(i).write(item);
+                        }
+                        iterated += 1;
+                    }
+                    None => break 'outer,
                 }
-            })
-            .count();
+            }
+        }
         // Safety: iterated slots have been initialized above
         unsafe { self.commit_unchecked(iterated) }
     }
@@ -713,14 +727,13 @@ impl<T> ReadChunk<'_, T> {
     }
 
     unsafe fn commit_unchecked(self, n: usize) -> usize {
-        for slot in self
-            .first_slice
+        self.first_slice
             .iter_mut()
             .chain(self.second_slice.iter_mut())
             .take(n)
-        {
-            slot.as_mut_ptr().drop_in_place();
-        }
+            .for_each(|slot| {
+                slot.as_mut_ptr().drop_in_place();
+            });
         let c = self.consumer;
         let head = c.buffer().increment(c.cached_head.get(), n);
         c.buffer().head.store(head, Ordering::Release);
@@ -751,12 +764,8 @@ impl<'a, T> IntoIterator for ReadChunk<'a, T> {
     /// Non-iterated items remain in the ring buffer.
     fn into_iter(self) -> Self::IntoIter {
         Self::IntoIter {
-            chunk_size: self.len(),
-            iter: self
-                .first_slice
-                .iter_mut()
-                .chain(self.second_slice.iter_mut()),
-            consumer: self.consumer,
+            chunk: self,
+            iterated: 0,
         }
     }
 }
@@ -770,22 +779,17 @@ impl<'a, T> IntoIterator for ReadChunk<'a, T> {
 /// Non-iterated items remain in the ring buffer.
 #[derive(Debug)]
 pub struct ReadChunkIntoIter<'a, T> {
-    chunk_size: usize,
-    iter: core::iter::Chain<
-        core::slice::IterMut<'a, MaybeUninit<T>>,
-        core::slice::IterMut<'a, MaybeUninit<T>>,
-    >,
-    consumer: &'a mut Consumer<T>,
+    chunk: ReadChunk<'a, T>,
+    iterated: usize,
 }
 
 impl<'a, T> Drop for ReadChunkIntoIter<'a, T> {
     /// Makes all iterated slots available for writing again.
     ///
     /// Non-iterated items remain in the ring buffer and are *not* dropped.
     fn drop(&mut self) {
-        let iterated = self.chunk_size - self.len();
-        let c = &self.consumer;
-        let head = c.buffer().increment(c.cached_head.get(), iterated);
+        let c = &self.chunk.consumer;
+        let head = c.buffer().increment(c.cached_head.get(), self.iterated);
         c.buffer().head.store(head, Ordering::Release);
         c.cached_head.set(head);
     }
@@ -795,11 +799,24 @@ impl<'a, T> Iterator for ReadChunkIntoIter<'a, T> {
     type Item = T;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.iter.next().map(|slot| unsafe { slot.as_ptr().read() })
+        self.chunk
+            .first_slice
+            .get(self.iterated)
+            .or_else(|| {
+                self.chunk
+                    .second_slice
+                    .get(self.iterated - self.chunk.first_slice.len())
+            })
+            .map(|slot| unsafe { slot.as_ptr().read() })
+            .map(|item| {
+                self.iterated += 1;
+                item
+            })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        self.iter.size_hint()
+        let size = self.chunk.first_slice.len() + self.chunk.second_slice.len() - self.iterated;
+        (size, Some(size))
     }
 }