From: Matt Corallo <git@bluematt.me>
Date: Sat, 4 Nov 2023 21:01:18 +0000 (+0000)
Subject: Pre-allocate the full `Vec` prior to serializing as a `Vec<u8>`
X-Git-Tag: v0.0.119~57^2~7
X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=commitdiff_plain;h=e4c6b70e8ee124ad91cbad84e36b5018da2d07a7;p=rust-lightning

Pre-allocate the full `Vec` prior to serializing as a `Vec<u8>`

We end up generating a substantial amount of allocations just
doubling `Vec`s when serializing to them, and our
`serialized_length` method is generally rather effecient, so we
just rely on it and allocate correctly up front.
---

diff --git a/lightning/src/util/ser.rs b/lightning/src/util/ser.rs
index 85b991c61..64d59da0d 100644
--- a/lightning/src/util/ser.rs
+++ b/lightning/src/util/ser.rs
@@ -199,8 +199,14 @@ pub trait Writeable {
 
 	/// Writes `self` out to a `Vec<u8>`.
 	fn encode(&self) -> Vec<u8> {
-		let mut msg = VecWriter(Vec::new());
+		let len = self.serialized_length();
+		let mut msg = VecWriter(Vec::with_capacity(len));
 		self.write(&mut msg).unwrap();
+		// Note that objects with interior mutability may change size between when we called
+		// serialized_length and when we called write. That's okay, but shouldn't happen during
+		// testing as most of our tests are not threaded.
+		#[cfg(test)]
+		debug_assert_eq!(len, msg.0.len());
 		msg.0
 	}
 
@@ -211,6 +217,7 @@ pub trait Writeable {
 		0u16.write(&mut msg).unwrap();
 		self.write(&mut msg).unwrap();
 		let len = msg.0.len();
+		debug_assert_eq!(len - 2, self.serialized_length());
 		msg.0[..2].copy_from_slice(&(len as u16 - 2).to_be_bytes());
 		msg.0
 	}