util: add Collected::to_bytes_mut

hyperium · Feb 8, 2024 · cc8f131 · cc8f131
1 parent 70ba87f
commit cc8f131
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/http-body-util/src/collected.rs b/http-body-util/src/collected.rs
@@ -4,7 +4,7 @@ use std::{
     task::{Context, Poll},
 };
 
-use bytes::{Buf, Bytes};
+use bytes::{Buf, Bytes, BytesMut};
 use http::HeaderMap;
 use http_body::{Body, Frame};
 
@@ -38,6 +38,11 @@ impl<B: Buf> Collected<B> {
         self.bufs.copy_to_bytes(self.bufs.remaining())
     }
 
+    /// Convert this body into a [`BytesMut`].
+    pub fn to_bytes_mut(mut self) -> BytesMut {
+        self.bufs.copy_to_bytes_mut(self.bufs.remaining())
+    }
+
     pub(crate) fn push_frame(&mut self, frame: Frame<B>) {
         let frame = match frame.into_data() {
             Ok(data) => {

diff --git a/http-body-util/src/util.rs b/http-body-util/src/util.rs
@@ -19,6 +19,14 @@ impl<T: Buf> BufList<T> {
     pub(crate) fn pop(&mut self) -> Option<T> {
         self.bufs.pop_front()
     }
+
+    #[inline]
+    pub(crate) fn copy_to_bytes_mut(&mut self, len: usize) -> BytesMut {
+        assert!(len <= self.remaining(), "`len` greater than remaining");
+        let mut bm = BytesMut::with_capacity(len);
+        bm.put(self.take(len));
+        bm
+    }
 }
 
 impl<T: Buf> Buf for BufList<T> {
@@ -77,10 +85,8 @@ impl<T: Buf> Buf for BufList<T> {
             }
             Some(front) if front.remaining() > len => front.copy_to_bytes(len),
             _ => {
-                assert!(len <= self.remaining(), "`len` greater than remaining");
-                let mut bm = BytesMut::with_capacity(len);
-                bm.put(self.take(len));
-                bm.freeze()
+                let bytes_mut = self.copy_to_bytes_mut(len);
+                bytes_mut.freeze()
             }
         }
     }