add mark/reset feature to ByteBuf and RingBuf

src/buf/byte.rs

@@ -13,7 +13,8 @@ pub struct ByteBuf {
     mem: alloc::MemRef,
     cap: u32,
     pos: u32,
-    lim: u32
+    lim: u32,
+    mark: Option<u32>,
 }
 
 impl ByteBuf {
@@ -35,6 +36,7 @@ impl ByteBuf {
             cap: 0,
             pos: 0,
             lim: 0,
+            mark: None,
         }
     }
 
@@ -46,6 +48,7 @@ impl ByteBuf {
             cap: cap,
             pos: pos,
             lim: lim,
+            mark: None,
         }
     }
 
@@ -70,7 +73,8 @@ impl ByteBuf {
             mem: mem,
             cap: capacity,
             pos: 0,
-            lim: capacity
+            lim: capacity,
+            mark: None,
         }
     }
 
@@ -112,6 +116,27 @@ impl ByteBuf {
         Bytes::of(self.to_seq_byte_str())
     }
 
+    /// Marks the current read location.
+    ///
+    /// Together with `reset`, this can be used to read from a section of the
+    /// buffer multiple times. The marked location will be cleared when the
+    /// buffer is flipped.
+    pub fn mark(&mut self) {
+        self.mark = Some(self.pos);
+    }
+
+    /// Resets the read position to the previously marked position.
+    ///
+    /// Together with `mark`, this can be used to read from a section of the
+    /// buffer multiple times.
+    ///
+    /// # Panics
+    ///
+    /// This method will panic if no mark has been set.
+    pub fn reset(&mut self) {
+        self.pos = self.mark.take().expect("no mark set");
+    }
+
     #[inline]
     fn pos(&self) -> usize {
         self.pos as usize
@@ -177,6 +202,26 @@ impl ROByteBuf {
     pub fn to_bytes(self) -> Bytes {
         self.buf.to_bytes()
     }
+
+    /// Marks the current read location.
+    ///
+    /// Together with `reset`, this can be used to read from a section of the
+    /// buffer multiple times.
+    pub fn mark(&mut self) {
+        self.buf.mark = Some(self.buf.pos);
+    }
+
+    /// Resets the read position to the previously marked position.
+    ///
+    /// Together with `mark`, this can be used to read from a section of the
+    /// buffer multiple times.
+    ///
+    /// # Panics
+    ///
+    /// This method will panic if no mark has been set.
+    pub fn reset(&mut self) {
+        self.buf.pos = self.buf.mark.take().expect("no mark set");
+    }
 }
 
 impl Buf for ROByteBuf {

src/buf/ring.rs

@@ -8,7 +8,8 @@ pub struct RingBuf {
     ptr: alloc::MemRef,  // Pointer to the memory
     cap: usize,          // Capacity of the buffer
     pos: usize,          // Offset of read cursor
-    len: usize           // Number of bytes to read
+    len: usize,          // Number of bytes to read
+    mark: Option<usize>, // Marked read position
 }
 
 // TODO: There are most likely many optimizations that can be made
@@ -20,7 +21,8 @@ impl RingBuf {
                 ptr: alloc::MemRef::none(),
                 cap: 0,
                 pos: 0,
-                len: 0
+                len: 0,
+                mark: None,
             }
         }
 
@@ -33,7 +35,8 @@ impl RingBuf {
             ptr: mem,
             cap: capacity,
             pos: 0,
-            len: 0
+            len: 0,
+            mark: None,
         }
     }
 
@@ -49,6 +52,29 @@ impl RingBuf {
         self.cap
     }
 
+    /// Marks the current read location.
+    ///
+    /// Together with `reset`, this can be used to read from a section of the
+    /// buffer multiple times. The mark will be cleared if it is overwritten
+    /// during a write.
+    pub fn mark(&mut self) {
+        self.mark = Some(self.pos);
+    }
+
+    /// Resets the read position to the previously marked position.
+    ///
+    /// Together with `mark`, this can be used to read from a section of the
+    /// buffer multiple times.
+    ///
+    /// # Panics
+    ///
+    /// This method will panic if no mark has been set,
+    pub fn reset(&mut self){
+        let mark = self.mark.take().expect("no mark set");
+        self.len = (self.len + self.pos + self.cap - mark) % self.cap;
+        self.pos = mark;
+    }
+
     fn read_remaining(&self) -> usize {
         self.len
     }
@@ -71,6 +97,11 @@ impl RingBuf {
     fn advance_writer(&mut self, mut cnt: usize) {
         cnt = cmp::min(cnt, self.write_remaining());
         self.len += cnt;
+        if let Some(mark) = self.mark {
+            if (self.pos + self.len) % self.cap > mark {
+                self.mark = None;
+            }
+        }
     }
 }
 

test/test_byte_buf.rs

@@ -30,6 +30,10 @@ pub fn test_byte_buf_read_write() {
     let mut buf = buf.flip();
     let mut dst = [0; 5];
 
+    buf.mark();
+    assert_eq!(5, buf.read(&mut dst[..]).unwrap());
+    assert_eq!(b"hello", &dst);
+    buf.reset();
     assert_eq!(5, buf.read(&mut dst[..]).unwrap());
     assert_eq!(b"hello", &dst);
 

test/test_ring.rs

@@ -19,9 +19,16 @@ pub fn test_initial_buf_empty() {
     assert_eq!(buf.bytes(), [1, 2, 3]);
 
     let mut out = [0u8; 3];
+
+    buf.mark();
     let bytes_read = buf.read(&mut out[..]).unwrap();;
     assert_eq!(bytes_read, 3);
     assert_eq!(out, [1, 2, 3]);
+    buf.reset();
+    let bytes_read = buf.read(&mut out[..]).unwrap();;
+    assert_eq!(bytes_read, 3);
+    assert_eq!(out, [1, 2, 3]);
+
     assert_eq!(MutBuf::remaining(&buf), 16);
     assert_eq!(Buf::remaining(&buf), 0);
 }
@@ -39,6 +46,11 @@ fn test_wrapping_write() {
     let bytes_written = buf.write(&[23;8][..]).unwrap();
     assert_eq!(bytes_written, 8);
 
+    buf.mark();
+    let bytes_read = buf.read(&mut out[..]).unwrap();
+    assert_eq!(bytes_read, 10);
+    assert_eq!(out, [42, 42, 23, 23, 23, 23, 23, 23, 23, 23]);
+    buf.reset();
     let bytes_read = buf.read(&mut out[..]).unwrap();
     assert_eq!(bytes_read, 10);
     assert_eq!(out, [42, 42, 23, 23, 23, 23, 23, 23, 23, 23]);
@@ -64,3 +76,16 @@ fn test_io_write_and_read() {
     assert_eq!(bytes_read, 8);
     assert_eq!(out, [2;8]);
 }
+
+#[test]
+#[should_panic]
+fn test_wrap_reset() {
+    use std::io::{Read, Write};
+
+    let mut buf = RingBuf::new(8);
+    buf.write(&[1, 2, 3, 4, 5, 6, 7]).unwrap();
+    buf.mark();
+    buf.read(&mut [0; 4]).unwrap();
+    buf.write(&[1, 2, 3, 4]).unwrap();
+    buf.reset();
+}