add get_array to Buf

Author: lperlaki

Cargo.toml

@@ -6,7 +6,7 @@ name = "bytes"
 # - Create "v1.x.y" git tag.
 version = "1.10.1"
 edition = "2018"
-rust-version = "1.39"
+rust-version = "1.51"
 license = "MIT"
 authors = [
     "Carl Lerche <me@carllerche.com>",

src/buf/buf_impl.rs

@@ -296,6 +296,30 @@ pub trait Buf {
             .unwrap_or_else(|error| panic_advance(&error));
     }
 
+    /// Gets a const known number of bytes from `self`
+    ///
+    /// The current position is advanced by `N`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::Buf;
+    ///
+    /// let mut buf = &b"hello world"[..];
+    /// assert_eq!(*b"hello", buf.get_array::<5>());
+    /// ```
+    ///
+    /// # Panics
+    ///
+    /// This function panics if there is not enough remaining data in `self`.
+    fn get_array<const N: usize>(&mut self) -> [u8; N]
+    where
+        Self: Sized,
+    {
+        self.try_get_array()
+            .unwrap_or_else(|error| panic_advance(&error))
+    }
+
     /// Gets an unsigned 8 bit integer from `self`.
     ///
     /// The current position is advanced by 1.
@@ -1178,6 +1202,61 @@ pub trait Buf {
         Ok(())
     }
 
+    /// Gets a const known number of bytes from `self`
+    ///
+    /// The current position is advanced by `N`.
+    ///
+    /// Returns `Err(TryGetError)` when there are not enough
+    /// remaining bytes to read the value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::Buf;
+    ///
+    /// let mut buf = &b"hello world"[..];
+    /// assert_eq!(Ok(*b"hello"), buf.try_get_array::<5>());
+    /// assert_eq!(6, buf.remaining());
+    /// ```
+    /// ```
+    /// use bytes::{Buf, TryGetError};
+    ///
+    /// let mut buf = &b"hel"[..];
+    /// assert_eq!(Err(TryGetError{requested: 5, available: 3}), buf.try_get_array::<5>());
+    /// assert_eq!(3, buf.remaining());
+    /// ```
+    ///
+    fn try_get_array<const N: usize>(&mut self) -> Result<[u8; N], TryGetError>
+    where
+        Self: Sized,
+    {
+        if self.remaining() < N {
+            return Err(TryGetError {
+                requested: N,
+                available: self.remaining(),
+            });
+        }
+
+        // try to convert directly from the bytes
+        // this Option<ret> trick is to avoid keeping a borrow on self
+        // when advance() is called (mut borrow) and to call bytes() only once
+        let ret = self
+            .chunk()
+            .get(..N)
+            .map(|src| unsafe { *(src as *const _ as *const [_; N]) });
+
+        if let Some(ret) = ret {
+            // if the direct conversion was possible, advance and return
+            self.advance(N);
+            return Ok(ret);
+        } else {
+            // if not we copy the bytes in a temp buffer then convert
+            let mut buf = [0; N];
+            self.copy_to_slice(&mut buf); // (do the advance)
+            return Ok(buf);
+        }
+    }
+
     /// Gets an unsigned 8 bit integer from `self`.
     ///
     /// The current position is advanced by 1.