Optimize length-prefixed writes

Author: devrandom

lightning/src/util/ser.rs

@@ -10,7 +10,7 @@
 //! A very simple serialization framework which is used to serialize/deserialize messages as well
 //! as ChannelsManagers and ChannelMonitors.
 
-use std::io::{Read, Write};
+use std::io::{Read, Write, Cursor};
 use std::collections::HashMap;
 use std::hash::Hash;
 use std::sync::Mutex;
@@ -62,6 +62,22 @@ impl Writer {
 	pub fn len(&self) -> usize {
 		self.0.len()
 	}
+
+	/// Insert the length of a data slice before the slice start
+	pub fn insert_length(&mut self, slice_start: usize, length_adjust: usize) -> Result<(), ::std::io::Error> {
+		let slice_end = self.len();
+		let length = slice_end - slice_start + length_adjust;
+		// Write the length after the slice, to check how many bytes
+		BigSize(length as u64).write(self)?;
+		let len_end = self.len();
+		let len_len = len_end - slice_end;
+		// Shift the slice to make space for the length
+		self.0.copy_within(slice_start..slice_end, slice_start + len_len);
+		let mut cursor = Cursor::new(&mut self.0);
+		cursor.set_position(slice_start as u64);
+		// Write the length again at the right spot
+		BigSize(length as u64).write_write(&mut cursor)
+	}
 }
 
 impl Write for Writer {
@@ -225,28 +241,34 @@ impl Readable for U48 {
 /// you're looking for an example of a variable-length integer to use for your own project, move
 /// along, this is a rather poor design.
 pub(crate) struct BigSize(pub u64);
-impl Writeable for BigSize {
-	#[inline]
-	fn write(&self, writer: &mut Writer) -> Result<(), ::std::io::Error> {
+impl BigSize {
+	pub fn write_write<W: Write>(&self, write: &mut W) -> Result<(), ::std::io::Error> {
 		match self.0 {
 			0...0xFC => {
-				(self.0 as u8).write(writer)
+				write.write_all(&[self.0 as u8])
 			},
 			0xFD...0xFFFF => {
-				0xFDu8.write(writer)?;
-				(self.0 as u16).write(writer)
+				write.write_all(&[0xFDu8])?;
+				write.write_all(&be16_to_array(self.0 as u16))
 			},
 			0x10000...0xFFFFFFFF => {
-				0xFEu8.write(writer)?;
-				(self.0 as u32).write(writer)
+				write.write_all(&[0xFEu8])?;
+				write.write_all(&be32_to_array(self.0 as u32))
 			},
 			_ => {
-				0xFFu8.write(writer)?;
-				(self.0 as u64).write(writer)
+				write.write_all(&[0xFFu8])?;
+				write.write_all(&be64_to_array(self.0 as u64))
 			},
 		}
 	}
 }
+
+impl Writeable for BigSize {
+	#[inline]
+	fn write(&self, writer: &mut Writer) -> Result<(), ::std::io::Error> {
+		self.write_write(writer)
+	}
+}
 impl Readable for BigSize {
 	#[inline]
 	fn read<R: Read>(reader: &mut R) -> Result<BigSize, DecodeError> {
@@ -609,10 +631,9 @@ impl<T: Writeable> Writeable for Option<T> {
 		match *self {
 			None => 0u8.write(w)?,
 			Some(ref data) => {
-				let mut len_calc = Writer::new();
-				data.write(&mut len_calc).expect("No in-memory data may fail to serialize");
-				BigSize(len_calc.len() as u64 + 1).write(w)?;
-				data.write(w)?;
+				let start = w.len();
+				data.write(w).expect("No in-memory data may fail to serialize");
+				w.insert_length(start, 1)?;
 			}
 		}
 		Ok(())

lightning/src/util/ser_macros.rs

@@ -9,13 +9,12 @@
 
 macro_rules! encode_tlv {
 	($stream: expr, {$(($type: expr, $field: expr)),*}) => { {
-		use util::ser::{BigSize, Writer};
+		use util::ser::BigSize;
 		$(
 			BigSize($type).write($stream)?;
-			let mut len_calc = Writer::new();
-			$field.write(&mut len_calc)?;
-			BigSize(len_calc.len() as u64).write($stream)?;
+			let start = $stream.len();
 			$field.write($stream)?;
+			$stream.insert_length(start, 0)?;
 		)*
 	} }
 }