[Codec] [Runtime] [p2p] Multi-buffer encoding for zero-copy of Bytes fields (#3491)

Co-authored-by: Patrick O'Grady <me@patrickogrady.xyz>

Author: danlaine

Cargo.lock

@@ -22,6 +22,14 @@ pub trait FixedSize {
 pub trait EncodeSize {
     /// Returns the encoded size of this value (in bytes).
     fn encode_size(&self) -> usize;
+
+    /// Returns the encoded size excluding bytes passed to [`BufsMut::push`]
+    /// during [`Write::write_bufs`]. Used to size the working buffer for inline
+    /// writes. Override alongside [`Write::write_bufs`] for types where large
+    /// [`Bytes`] fields go via push; failing to do so will over-allocate.
+    fn encode_inline_size(&self) -> usize {
+        self.encode_size()
+    }
 }
 
 // Automatically implement `EncodeSize` for types that are `FixedSize`.
@@ -37,6 +45,13 @@ pub trait Write {
     ///
     /// Implementations should panic if the buffer doesn't have enough capacity.
     fn write(&self, buf: &mut impl BufMut);
+
+    /// Writes to a [`BufsMut`], allowing existing [`Bytes`] chunks to be
+    /// appended via [`BufsMut::push`] instead of written inline. Must encode
+    /// to the same format as [`Write::write`]. Defaults to [`Write::write`].
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.write(buf);
+    }
 }
 
 /// Trait for types that can be read (decoded) from a byte buffer.
@@ -194,6 +209,13 @@ pub trait CodecFixedShared: CodecFixed<Cfg = ()> + Send + Sync {}
 // Automatically implement `CodecFixedShared` for types that meet all bounds.
 impl<T: CodecFixed<Cfg = ()> + Send + Sync> CodecFixedShared for T {}
 
+/// A [`BufMut`] that can also append pre-existing [`Bytes`] chunks.
+pub trait BufsMut: BufMut {
+    /// Appends a [`Bytes`] chunk instead of writing its contents inline into
+    /// the destination buffer.
+    fn push(&mut self, bytes: impl Into<Bytes>);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

codec/src/codec.rs

@@ -6,7 +6,7 @@
 extern crate alloc;
 
 use crate::{
-    codec::{EncodeSize, Read, Write},
+    codec::{BufsMut, EncodeSize, Read, Write},
     error::Error,
     types::read_ordered_map,
     RangeCfg,
@@ -28,6 +28,16 @@ impl<K: Ord + Eq + Write, V: Write> Write for BTreeMap<K, V> {
             v.write(buf);
         }
     }
+
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.len().write(buf);
+
+        // Keys are already sorted in BTreeMap, so we can iterate directly
+        for (k, v) in self {
+            k.write_bufs(buf);
+            v.write_bufs(buf);
+        }
+    }
 }
 
 impl<K: Ord + Eq + EncodeSize, V: EncodeSize> EncodeSize for BTreeMap<K, V> {
@@ -42,6 +52,18 @@ impl<K: Ord + Eq + EncodeSize, V: EncodeSize> EncodeSize for BTreeMap<K, V> {
         }
         size
     }
+
+    fn encode_inline_size(&self) -> usize {
+        // Start with the size of the length prefix
+        let mut size = self.len().encode_size();
+
+        // Add the encoded size of each key and value
+        for (k, v) in self {
+            size += k.encode_inline_size();
+            size += v.encode_inline_size();
+        }
+        size
+    }
 }
 
 impl<K: Read + Clone + Ord + Eq, V: Read + Clone> Read for BTreeMap<K, V> {

codec/src/types/btree_map.rs

@@ -6,7 +6,7 @@
 extern crate alloc;
 
 use crate::{
-    codec::{EncodeSize, Read, Write},
+    codec::{BufsMut, EncodeSize, Read, Write},
     error::Error,
     types::read_ordered_set,
     RangeCfg,
@@ -27,6 +27,15 @@ impl<K: Ord + Eq + Write> Write for BTreeSet<K> {
             item.write(buf);
         }
     }
+
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.len().write(buf);
+
+        // Items are already sorted in BTreeSet, so we can iterate directly
+        for item in self {
+            item.write_bufs(buf);
+        }
+    }
 }
 
 impl<K: Ord + Eq + EncodeSize> EncodeSize for BTreeSet<K> {
@@ -37,6 +46,14 @@ impl<K: Ord + Eq + EncodeSize> EncodeSize for BTreeSet<K> {
         }
         size
     }
+
+    fn encode_inline_size(&self) -> usize {
+        let mut size = self.len().encode_size();
+        for item in self {
+            size += item.encode_inline_size();
+        }
+        size
+    }
 }
 
 impl<K: Read + Clone + Ord + Eq> Read for BTreeSet<K> {

codec/src/types/btree_set.rs

@@ -3,7 +3,7 @@
 //! For portability and consistency between architectures,
 //! the length of the [Bytes] must fit within a [u32].
 
-use crate::{util::at_least, EncodeSize, Error, RangeCfg, Read, Write};
+use crate::{util::at_least, BufsMut, EncodeSize, Error, RangeCfg, Read, Write};
 use bytes::{Buf, BufMut, Bytes};
 
 impl Write for Bytes {
@@ -12,13 +12,24 @@ impl Write for Bytes {
         self.len().write(buf);
         buf.put_slice(self);
     }
+
+    #[inline]
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.len().write(buf);
+        buf.push(self.clone());
+    }
 }
 
 impl EncodeSize for Bytes {
     #[inline]
     fn encode_size(&self) -> usize {
         self.len().encode_size() + self.len()
     }
+
+    #[inline]
+    fn encode_inline_size(&self) -> usize {
+        self.len().encode_size()
+    }
 }
 
 impl Read for Bytes {

codec/src/types/bytes.rs

@@ -4,7 +4,7 @@
 //! the size of the map must fit within a [u32].
 
 use crate::{
-    codec::{EncodeSize, Read, Write},
+    codec::{BufsMut, EncodeSize, Read, Write},
     error::Error,
     types::read_ordered_map,
     RangeCfg,
@@ -28,6 +28,18 @@ impl<K: Ord + Hash + Eq + Write, V: Write> Write for HashMap<K, V> {
             v.write(buf);
         }
     }
+
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.len().write(buf);
+
+        // Sort the keys to ensure deterministic encoding
+        let mut entries: Vec<_> = self.iter().collect();
+        entries.sort_by(|a, b| a.0.cmp(b.0));
+        for (k, v) in entries {
+            k.write_bufs(buf);
+            v.write_bufs(buf);
+        }
+    }
 }
 
 impl<K: Ord + Hash + Eq + EncodeSize, V: EncodeSize> EncodeSize for HashMap<K, V> {
@@ -43,6 +55,19 @@ impl<K: Ord + Hash + Eq + EncodeSize, V: EncodeSize> EncodeSize for HashMap<K, V
         }
         size
     }
+
+    fn encode_inline_size(&self) -> usize {
+        // Start with the size of the length prefix
+        let mut size = self.len().encode_size();
+
+        // Add the encoded size of each key and value
+        // Note: Iteration order doesn't matter for size calculation.
+        for (k, v) in self {
+            size += k.encode_inline_size();
+            size += v.encode_inline_size();
+        }
+        size
+    }
 }
 
 // Read implementation for HashMap

codec/src/types/hash_map.rs

@@ -4,7 +4,7 @@
 //! the size of the set must fit within a [u32].
 
 use crate::{
-    codec::{EncodeSize, Read, Write},
+    codec::{BufsMut, EncodeSize, Read, Write},
     error::Error,
     types::read_ordered_set,
     RangeCfg,
@@ -25,6 +25,17 @@ impl<K: Ord + Hash + Eq + Write> Write for HashSet<K> {
             item.write(buf);
         }
     }
+
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.len().write(buf);
+
+        // Sort the items to ensure deterministic encoding
+        let mut items: Vec<_> = self.iter().collect();
+        items.sort();
+        for item in items {
+            item.write_bufs(buf);
+        }
+    }
 }
 
 impl<K: Ord + Hash + Eq + EncodeSize> EncodeSize for HashSet<K> {
@@ -37,6 +48,16 @@ impl<K: Ord + Hash + Eq + EncodeSize> EncodeSize for HashSet<K> {
         }
         size
     }
+
+    fn encode_inline_size(&self) -> usize {
+        let mut size = self.len().encode_size();
+
+        // Note: Iteration order doesn't matter for size calculation.
+        for item in self {
+            size += item.encode_inline_size();
+        }
+        size
+    }
 }
 
 impl<K: Read + Clone + Ord + Hash + Eq> Read for HashSet<K> {

codec/src/types/hash_set.rs

@@ -1,6 +1,6 @@
 //! This module exports the [`Lazy`] type.
 
-use crate::{Decode, Encode, EncodeSize, FixedSize, Read, Write};
+use crate::{BufsMut, Decode, Encode, EncodeSize, FixedSize, Read, Write};
 use bytes::{Buf, Bytes};
 use core::hash::Hash;
 #[cfg(feature = "std")]
@@ -160,6 +160,15 @@ impl<T: Read + EncodeSize> EncodeSize for Lazy<T> {
             .expect("Lazy should have a value if pending is None")
             .encode_size()
     }
+
+    fn encode_inline_size(&self) -> usize {
+        if self.pending.is_some() {
+            return 0;
+        }
+        self.get()
+            .expect("Lazy should have a value if pending is None")
+            .encode_inline_size()
+    }
 }
 
 impl<T: Read + Write> Write for Lazy<T> {
@@ -173,6 +182,17 @@ impl<T: Read + Write> Write for Lazy<T> {
             .expect("Lazy should have a value if pending is None")
             .write(buf);
     }
+
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        if let Some(pending) = &self.pending {
+            // Write raw bytes without length prefix (Bytes::write_bufs adds a length prefix)
+            buf.push(pending.bytes.clone());
+            return;
+        }
+        self.get()
+            .expect("Lazy should have a value if pending is None")
+            .write_bufs(buf);
+    }
 }
 
 impl<T: Read + FixedSize> Read for Lazy<T> {

codec/src/types/lazy.rs

@@ -19,7 +19,8 @@
 //!   endian ambiguity.
 
 use crate::{
-    util::at_least, varint::UInt, EncodeSize, Error, FixedSize, RangeCfg, Read, ReadExt, Write,
+    util::at_least, varint::UInt, BufsMut, EncodeSize, Error, FixedSize, RangeCfg, Read, ReadExt,
+    Write,
 };
 use bytes::{Buf, BufMut};
 use core::num::{NonZeroU16, NonZeroU32, NonZeroU64};
@@ -196,13 +197,27 @@ impl<T: Write> Write for Option<T> {
             inner.write(buf);
         }
     }
+
+    #[inline]
+    fn write_bufs(&self, buf: &mut impl BufsMut) {
+        self.is_some().write(buf);
+        if let Some(inner) = self {
+            inner.write_bufs(buf);
+        }
+    }
 }
 
 impl<T: EncodeSize> EncodeSize for Option<T> {
     #[inline]
     fn encode_size(&self) -> usize {
         self.as_ref().map_or(1, |inner| 1 + inner.encode_size())
     }
+
+    #[inline]
+    fn encode_inline_size(&self) -> usize {
+        self.as_ref()
+            .map_or(1, |inner| 1 + inner.encode_inline_size())
+    }
 }
 
 impl<T: Read> Read for Option<T> {