From d8cb390d606638e0bb24208a56996b680e9d21b4 Mon Sep 17 00:00:00 2001
From: Ruben De Smet <ruben.de.smet@rubdos.be>
Date: Tue, 6 Jul 2021 16:00:10 +0200
Subject: [PATCH] Add encode to Vec<u8> (#378)

* Message::{encode_to_vec,encode_length_delimited_to_vec}

* test for Message::encode_to_vec

Co-authored-by: Lucio Franco <luciofranco14@gmail.com>
---
 src/message.rs   | 26 ++++++++++++++++++++++++++
 tests/src/lib.rs |  7 +++++++
 2 files changed, 33 insertions(+)

diff --git a/src/message.rs b/src/message.rs
index 4706387e6..112c7b89f 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -59,6 +59,18 @@ pub trait Message: Debug + Send + Sync {
         Ok(())
     }
 
+    #[cfg(feature = "std")]
+    /// Encodes the message to a newly allocated buffer.
+    fn encode_to_vec(&self) -> Vec<u8>
+    where
+        Self: Sized,
+    {
+        let mut buf = Vec::with_capacity(self.encoded_len());
+
+        self.encode_raw(&mut buf);
+        buf
+    }
+
     /// Encodes the message with a length-delimiter to a buffer.
     ///
     /// An error will be returned if the buffer does not have sufficient capacity.
@@ -78,6 +90,20 @@ pub trait Message: Debug + Send + Sync {
         Ok(())
     }
 
+    #[cfg(feature = "std")]
+    /// Encodes the message with a length-delimiter to a newly allocated buffer.
+    fn encode_length_delimited_to_vec(&self) -> Vec<u8>
+    where
+        Self: Sized,
+    {
+        let len = self.encoded_len();
+        let mut buf = Vec::with_capacity(len + encoded_len_varint(len as u64));
+
+        encode_varint(len as u64, &mut buf);
+        self.encode_raw(&mut buf);
+        buf
+    }
+
     /// Decodes an instance of the message from a buffer.
     ///
     /// The entire buffer will be consumed.
diff --git a/tests/src/lib.rs b/tests/src/lib.rs
index 4b13ed583..205205805 100644
--- a/tests/src/lib.rs
+++ b/tests/src/lib.rs
@@ -180,6 +180,7 @@ where
     if let Err(error) = roundtrip.encode(&mut buf2) {
         return RoundtripResult::Error(error.into());
     }
+    let buf3 = roundtrip.encode_to_vec();
 
     /*
     // Useful for debugging:
@@ -192,6 +193,12 @@ where
         return RoundtripResult::Error(anyhow!("roundtripped encoded buffers do not match"));
     }
 
+    if buf1 != buf3 {
+        return RoundtripResult::Error(anyhow!(
+            "roundtripped encoded buffers do not match with `encode_to_vec`"
+        ));
+    }
+
     RoundtripResult::Ok(buf1)
 }