Don't clone in non-contiguous case

Author: fuzzypixelz

commons/zenoh-buffers/src/zbuf.rs

@@ -12,7 +12,12 @@
 //   ZettaScale Zenoh Team, <[email protected]>
 //
 use alloc::{sync::Arc, vec::Vec};
-use core::{cmp, iter, num::NonZeroUsize, ptr::NonNull};
+use core::{
+    cmp::{self, Ordering},
+    iter,
+    num::NonZeroUsize,
+    ptr::NonNull,
+};
 #[cfg(feature = "std")]
 use std::io;
 
@@ -82,6 +87,17 @@ impl ZBuf {
     fn opt_zslice_writer(&mut self) -> Option<ZSliceWriter> {
         self.slices.last_mut().and_then(|s| s.writer())
     }
+
+    pub fn chunks(&self, chunk_size: usize) -> ZBufChunkIterator<'_> {
+        assert_ne!(chunk_size, 0, "cannot split ZBuf into chunks of size 0");
+
+        ZBufChunkIterator {
+            chunk_size,
+            buf: self,
+            slice_cursor: 0,
+            byte_cursor: 0,
+        }
+    }
 }
 
 // Buffer
@@ -669,3 +685,97 @@ mod tests {
         reader.seek(std::io::SeekFrom::Current(-100)).unwrap_err();
     }
 }
+
+pub struct ZBufChunkIterator<'a> {
+    chunk_size: usize,
+    buf: &'a ZBuf,
+    slice_cursor: usize,
+    byte_cursor: usize,
+}
+
+impl ZBufChunkIterator<'_> {
+    fn inc_slice_cursor(&mut self) {
+        self.slice_cursor += 1;
+        self.byte_cursor = 0;
+    }
+
+    fn inc_byte_cursor(&mut self, mut amt: usize) {
+        while let Some(slice) = self.buf.slices.get(self.slice_cursor) {
+            let rem = slice.len() - self.byte_cursor;
+
+            match rem.cmp(&amt) {
+                Ordering::Less => {
+                    self.inc_slice_cursor();
+                    amt -= rem;
+                    continue;
+                }
+                Ordering::Equal => {
+                    self.inc_slice_cursor();
+                    break;
+                }
+                Ordering::Greater => {
+                    self.byte_cursor += amt;
+                    break;
+                }
+            }
+        }
+    }
+}
+
+impl Iterator for ZBufChunkIterator<'_> {
+    type Item = ZBuf;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let slice = self.buf.slices.get(self.slice_cursor)?;
+
+        let rem = slice.len() - self.byte_cursor;
+
+        match rem.cmp(&self.chunk_size) {
+            Ordering::Less => {
+                let mut buf = ZBuf::empty();
+                let subslice = slice.subslice(self.byte_cursor..).unwrap();
+                buf.push_zslice(subslice);
+                self.inc_slice_cursor();
+
+                while let Some(slice) = self.buf.slices.get(self.slice_cursor) {
+                    let missing = self.chunk_size - buf.len();
+
+                    match slice.len().cmp(&missing) {
+                        Ordering::Less => {
+                            buf.push_zslice(slice.clone());
+                            self.inc_slice_cursor();
+                            continue;
+                        }
+                        Ordering::Equal => {
+                            buf.push_zslice(slice.clone());
+                            self.inc_slice_cursor();
+                            break;
+                        }
+                        Ordering::Greater => {
+                            let subslice = slice.subslice(..missing).unwrap();
+                            buf.push_zslice(subslice);
+                            self.inc_byte_cursor(missing);
+                            break;
+                        }
+                    }
+                }
+
+                Some(buf)
+            }
+            Ordering::Equal => {
+                let subslice = slice.subslice(self.byte_cursor..).unwrap();
+
+                self.inc_slice_cursor();
+                Some(ZBuf::from(subslice))
+            }
+            Ordering::Greater => {
+                let subslice = slice
+                    .subslice(self.byte_cursor..self.byte_cursor + self.chunk_size)
+                    .unwrap();
+
+                self.inc_byte_cursor(self.chunk_size);
+                Some(ZBuf::from(subslice))
+            }
+        }
+    }
+}

commons/zenoh-buffers/src/zslice.rs

@@ -198,22 +198,6 @@ impl ZSlice {
             None
         }
     }
-
-    pub fn chunks(self, chunk_size: usize) -> impl Iterator<Item = Self> {
-        assert_ne!(chunk_size, 0, "cannot split ZSlice into chunks of size 0");
-
-        let n = self.len() / chunk_size;
-        let rem = self.len() % chunk_size;
-        let n = if rem > 0 { n + 1 } else { n };
-
-        (0..n).map(move |i| ZSlice {
-            buf: self.buf.clone(),
-            start: self.start + i * chunk_size,
-            end: usize::min(self.end, self.start + (i + 1) * chunk_size),
-            #[cfg(feature = "shared-memory")]
-            kind: ZSliceKind::Raw,
-        })
-    }
 }
 
 impl Deref for ZSlice {
@@ -442,58 +426,3 @@ impl ZSlice {
         (0..len).map(|_| rng.gen()).collect::<Vec<u8>>().into()
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn zslice() {
-        let buf = crate::vec::uninit(16);
-        let mut zslice: ZSlice = buf.clone().into();
-        assert_eq!(buf.as_slice(), zslice.as_slice());
-
-        // SAFETY: buffer slize size is not modified
-        let mut_slice = unsafe { zslice.downcast_mut::<Vec<u8>>() }.unwrap();
-
-        mut_slice[..buf.len()].clone_from_slice(&buf[..]);
-
-        assert_eq!(buf.as_slice(), zslice.as_slice());
-    }
-
-    #[test]
-    fn test_chunks_inexact_division() {
-        let vec = (0..15).collect::<Vec<_>>();
-        let zslice = ZSlice::from(vec);
-
-        const EXPECTED_CHUNKS: &[&[u8]] =
-            &[&[0, 1, 2, 3], &[4, 5, 6, 7], &[8, 9, 10, 11], &[12, 13, 14]];
-        const CHUNK_SIZE: usize = 4;
-
-        for (found, expected) in zslice.chunks(CHUNK_SIZE).zip(EXPECTED_CHUNKS) {
-            assert_eq!(found.as_slice(), *expected);
-            assert!(found.len() <= CHUNK_SIZE);
-        }
-    }
-
-    #[test]
-    fn test_chunks_exact_division() {
-        let vec = (0..8).collect::<Vec<_>>();
-        let zslice = ZSlice::from(vec);
-
-        const EXPECTED_CHUNKS: &[&[u8]] = &[&[0, 1], &[2, 3], &[4, 5], &[6, 7]];
-        const CHUNK_SIZE: usize = 2;
-
-        for (found, expected) in zslice.chunks(CHUNK_SIZE).zip(EXPECTED_CHUNKS) {
-            assert_eq!(found.as_slice(), *expected);
-            assert!(found.len() <= CHUNK_SIZE);
-        }
-    }
-
-    #[test]
-    fn test_chunks_empty() {
-        let zslice = ZSlice::empty();
-
-        assert_eq!(zslice.chunks(3).next(), None);
-    }
-}

zenoh/src/api/bytes.rs

@@ -18,7 +18,7 @@ use std::{borrow::Cow, fmt::Debug, mem, str::Utf8Error};
 use zenoh_buffers::{
     buffer::{Buffer, SplitBuffer},
     reader::{HasReader, Reader},
-    ZBuf, ZBufReader, ZSlice, ZSliceBuffer,
+    ZBuf, ZBufChunkIterator, ZBufReader, ZSlice, ZSliceBuffer,
 };
 use zenoh_protocol::zenoh::ext::AttachmentType;
 
@@ -213,7 +213,7 @@ impl ZBytes {
 
     /// Return an iterator of chunks in [`ZBytes`].
     ///
-    /// This method does not entail cloning the underlying buffer unless [`ZBytes`] is not
+    /// This method does not entail cloning the underlying buffer even if [`ZBytes`] is not
     /// contiguous.
     ///
     /// All chunks have size less then or equal to `chunk_size`.
@@ -222,11 +222,8 @@ impl ZBytes {
     ///
     /// Panics if `chunk_size` is zero.
     #[zenoh_macros::unstable]
-    pub fn chunks(&self, chunk_size: usize) -> impl Iterator<Item = ZBytes> {
-        self.0
-            .to_zslice()
-            .chunks(chunk_size)
-            .map(|zs| ZBytes::from(ZBuf::from(zs)))
+    pub fn chunks(&self, chunk_size: usize) -> ZBytesChunkIterator {
+        ZBytesChunkIterator(self.0.chunks(chunk_size))
     }
 }
 
@@ -507,3 +504,16 @@ impl<const ID: u8> From<AttachmentType<ID>> for ZBytes {
         this.buffer.into()
     }
 }
+
+/// An iterator of chunks in [`ZBytes`], see [`ZBytes::chunks`].
+#[zenoh_macros::unstable]
+pub struct ZBytesChunkIterator<'a>(ZBufChunkIterator<'a>);
+
+#[zenoh_macros::unstable]
+impl Iterator for ZBytesChunkIterator<'_> {
+    type Item = ZBytes;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next().map(ZBytes::from)
+    }
+}