impl Into<Vec<u8>> for EasyBuf, non-allocating if possible (cf. tokio…

…-rs#120)
manuel-woelker · Feb 8, 2017 · e0027b2 · e0027b2
1 parent 062d3fc
commit e0027b2
Showing 1 changed file with 37 additions and 2 deletions.
diff --git a/src/io/frame.rs b/src/io/frame.rs
@@ -209,8 +209,20 @@ impl fmt::Debug for EasyBuf {
 }
 
 impl Into<Vec<u8>> for EasyBuf {
-    fn into(self) -> Vec<u8> {
-        self.as_ref().to_vec()
+    fn into(mut self) -> Vec<u8> {
+        match Arc::try_unwrap(self.buf) {
+            Ok(mut v) => {
+                // Unshared, reuse the Vec
+                v.truncate(self.end);
+                v.drain(..self.start);
+                v
+            },
+            Err(buf) => {
+                // Shared, move buf back into place
+                self.buf = buf;
+                self.as_ref().to_vec()
+            }
+        }
     }
 }
 
@@ -453,20 +465,43 @@ mod tests {
         let vec: Vec<u8> = (0u8..10u8).collect();
         let reference = vec.clone();
         let buf: EasyBuf = vec.into();
+        let original_pointer = buf.buf.as_ref().as_ptr();
         let result: Vec<u8> = buf.into();
         assert_eq!(result, reference);
+        let new_pointer = result.as_ptr();
+        assert_eq!(original_pointer, new_pointer, "Into<Vec<u8>> should be reuse the exclusive Vec");
     }
 
     #[test]
     fn easybuf_into_vec_sliced() {
         let vec: Vec<u8> = (0u8..10u8).collect();
         let mut buf: EasyBuf = vec.into();
+        let original_pointer = buf.buf.as_ref().as_ptr();
         buf.split_off(9);
         buf.drain_to(3);
         let result: Vec<u8> = buf.into();
         let reference: Vec<u8> = (3u8..9u8).collect();
         assert_eq!(result, reference);
+        let new_pointer = result.as_ptr();
+        assert_eq!(original_pointer, new_pointer, "Into<Vec<u8>> should be reuse the exclusive Vec");
     }
 
+    #[test]
+    fn easybuf_into_vec_sliced_allocating() {
+        let vec: Vec<u8> = (0u8..10u8).collect();
+        let mut buf: EasyBuf = vec.into();
+        let original_pointer = buf.buf.as_ref().as_ptr();
+        // Create a clone to create second reference to this EasyBuf and force allocation
+        let original = buf.clone();
+        buf.split_off(9);
+        buf.drain_to(3);
+        let result: Vec<u8> = buf.into();
+        let reference: Vec<u8> = (3u8..9u8).collect();
+        assert_eq!(result, reference);
+        let original_reference: EasyBuf =(0u8..10u8).collect::<Vec<u8>>().into();
+        assert_eq!(original.as_ref(), original_reference.as_ref());
+        let new_pointer = result.as_ptr();
+        assert_ne!(original_pointer, new_pointer, "A new vec should be allocated");
+    }
 
 }