Optimize length prefix writing

Author: devrandom

lightning/src/util/ser.rs

@@ -63,22 +63,29 @@ impl Writer {
 		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> {
+	/// Write the data with a varint length prefix.
+	/// Use length_adjust to add to the length for special cases, such as writing an Option.
+	pub fn write_all_with_length<T: Writeable>(&mut self, data: &T, length_adjust: usize) -> Result<(), ::std::io::Error> {
+		// Most common case is a single byte length prefix, so leave a space for that
+		self.0.push(0);
+		let start = self.len();
+		data.write(self).expect("No in-memory data may fail to serialize");
+		self.insert_length(start, length_adjust)
+	}
+
+	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;
-		self.0.resize(slice_end, 0);
-		// Shift the slice to make space for the length
-		for _ in 0 .. len_len {
+		let big_size = BigSize(length as u64);
+		let len_len = big_size.expected_length();
+		// At this point we have a placeholder for a varint of length 1 just before the start of data
+		// Shift the slice to make space for the length.  If len_len is 1, no shift is needed.
+		for _ in 1..len_len {
 			self.0.insert(slice_start, 0);
 		}
 		let mut cursor = Cursor::new(&mut self.0);
-		cursor.set_position(slice_start as u64);
-		// Write the length again at the right spot
+		cursor.set_position((slice_start - 1) as u64);
+		// Write the length in the space we created before the start of data
 		BigSize(length as u64).write_write(&mut cursor)
 	}
 }
@@ -264,6 +271,15 @@ impl BigSize {
 			},
 		}
 	}
+
+	pub fn expected_length(&self) -> usize {
+		match self.0 {
+			0...0xFC => 1,
+			0xFD...0xFFFF => 3,
+			0x10000...0xFFFFFFFF => 5,
+			_ => 9,
+		}
+	}
 }
 
 impl Writeable for BigSize {
@@ -634,9 +650,7 @@ impl<T: Writeable> Writeable for Option<T> {
 		match *self {
 			None => 0u8.write(w)?,
 			Some(ref data) => {
-				let start = w.len();
-				data.write(w).expect("No in-memory data may fail to serialize");
-				w.insert_length(start, 1)?;
+				w.write_all_with_length(data, 1)?;
 			}
 		}
 		Ok(())
@@ -756,3 +770,56 @@ impl<A: Writeable, B: Writeable> Writeable for (A, B) {
 		self.1.write(w)
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use std::io::Cursor;
+	use util::ser::*;
+
+	#[test]
+	fn big_size() {
+		for i in vec![0, 0xfc, 0xfd, 0xffff, 0x10000, 0xFFFFFFFF, 0xFFFFFFFFF] {
+			let mut writer = Writer::new();
+			let size = BigSize(i);
+			size.write(&mut writer).unwrap();
+			assert_eq!(writer.len(), size.expected_length());
+			let read_size = BigSize::read(&mut Cursor::new(writer.0)).unwrap();
+			assert_eq!(read_size.0, size.0);
+		}
+	}
+
+	#[test]
+	fn length_prefix_zero() {
+		impl_array!(0);
+		let mut writer = Writer::new();
+		writer.write_all_with_length(&[], 0).unwrap();
+		// Writing a zero-length length-prefixed item should result in just one zero byte
+		assert_eq!(writer.0, vec![0]);
+	}
+
+	#[test]
+	fn length_prefix_one() {
+		impl_array!(1);
+		let mut writer = Writer::new();
+		writer.write_all_with_length(&[0x11u8], 0).unwrap();
+		assert_eq!(writer.0, vec![1, 0x11]);
+	}
+
+	#[test]
+	fn length_prefix() {
+		for size in vec![3, 0xfc, 0xfd, 0xffff, 0x10000] {
+			let mut writer = Writer::new();
+			let mut data = Vec::new();
+			// A Vec serializes its own 2 bytes length, so adjust this
+			data.resize(size-3, 0x11u8);
+			// Writing an Option is convenient for round-trip checking of length-prefixed writes
+			Some(data.clone()).write(&mut writer).unwrap();
+			let mut reader = Cursor::new(&writer.0);
+			let read_data: Option<Vec<u8>> = Readable::read(&mut reader).unwrap();
+            assert_eq!(data, read_data.unwrap());
+			// Ensure everything was consumed
+            assert_eq!(reader.position() as usize, writer.len());
+			println!("{}", writer.len())
+		}
+	}
+}

lightning/src/util/ser_macros.rs

@@ -12,9 +12,7 @@ macro_rules! encode_tlv {
 		use util::ser::BigSize;
 		$(
 			BigSize($type).write($stream)?;
-			let start = $stream.len();
-			$field.write($stream)?;
-			$stream.insert_length(start, 0)?;
+			$stream.write_all_with_length(&$field, 0)?;
 		)*
 	} }
 }