Merge pull request #81 from ralexstokes/refactor-ser

harden typing to increase safety of serialization routine

ssz-rs-derive/src/lib.rs

@@ -88,39 +88,20 @@ fn derive_serialize_impl(data: &Data) -> TokenStream {
                     "this type of struct is currently not supported by this derive macro"
                 ),
             };
-            let serialization_by_field = fields.iter().map(|f| {
-                let field_type = &f.ty;
-                match &f.ident {
-                    Some(field_name) => quote_spanned! { f.span() =>
-                        let mut element_buffer = Vec::with_capacity(<#field_type>::size_hint());
-                        self.#field_name.serialize(&mut element_buffer)?;
-
-                        let element_buffer_len = element_buffer.len();
-                        if <#field_type>::is_variable_size() {
-                            fixed.push(None);
-                            fixed_lengths_sum += #BYTES_PER_LENGTH_OFFSET;
-                            variable.append(&mut element_buffer);
-                            variable_lengths.push(element_buffer_len);
-                        } else {
-                            fixed.push(Some(element_buffer));
-                            fixed_lengths_sum += element_buffer_len;
-                            variable_lengths.push(0)
-                        }
-                    },
-                    None => panic!("should have already returned an impl"),
-                }
+            let serialization_by_field = fields.iter().map(|f| match &f.ident {
+                Some(field_name) => quote_spanned! { f.span() =>
+                    serializer.with_element(&self.#field_name)?;
+                },
+                None => panic!("should have already returned an impl"),
             });
 
             quote! {
                 fn serialize(&self, buffer: &mut Vec<u8>) -> Result<usize, ssz_rs::SerializeError> {
-                    let mut fixed = vec![];
-                    let mut variable = vec![];
-                    let mut variable_lengths = vec![];
-                    let mut fixed_lengths_sum = 0;
+                    let mut serializer = ssz_rs::__internal::Serializer::default();
 
                     #(#serialization_by_field)*
 
-                    ssz_rs::__internal::serialize_composite_from_components(fixed, variable, variable_lengths, fixed_lengths_sum, buffer)
+                    serializer.serialize(buffer)
                 }
             }
         }

ssz-rs/src/array.rs

@@ -8,7 +8,7 @@ use crate::{
     error::{InstanceError, TypeError},
     lib::*,
     merkleization::{elements_to_chunks, merkleize, pack, MerkleizationError, Merkleized, Node},
-    ser::{serialize_composite, Serialize, SerializeError},
+    ser::{Serialize, SerializeError, Serializer},
     SimpleSerialize, Sized,
 };
 
@@ -35,7 +35,11 @@ macro_rules! define_ssz_for_array_of_size {
                 if $n == 0 {
                     return Err(TypeError::InvalidBound($n).into())
                 }
-                serialize_composite(self, buffer)
+                let mut serializer = Serializer::default();
+                for element in self {
+                    serializer.with_element(element)?;
+                }
+                serializer.serialize(buffer)
             }
         }
 

ssz-rs/src/lib.rs

@@ -135,6 +135,6 @@ pub mod __internal {
     // exported for derive macro to avoid code duplication...
     pub use crate::{
         merkleization::{merkleize, mix_in_selector},
-        ser::serialize_composite_from_components,
+        ser::Serializer,
     };
 }

ssz-rs/src/list.rs

@@ -5,7 +5,7 @@ use crate::{
     merkleization::{
         elements_to_chunks, merkleize, mix_in_length, pack, MerkleizationError, Merkleized, Node,
     },
-    ser::{serialize_composite, Serialize, SerializeError},
+    ser::{Serialize, SerializeError, Serializer},
     SimpleSerialize, Sized,
 };
 #[cfg(feature = "serde")]
@@ -177,7 +177,11 @@ where
         if self.len() > N {
             return Err(InstanceError::Bounded { bound: N, provided: self.len() }.into())
         }
-        serialize_composite(&self.data, buffer)
+        let mut serializer = Serializer::default();
+        for element in &self.data {
+            serializer.with_element(element)?;
+        }
+        serializer.serialize(buffer)
     }
 }
 

ssz-rs/src/ser.rs

@@ -11,8 +11,8 @@ const MAXIMUM_LENGTH: u64 = 2u64.pow((8 * BYTES_PER_LENGTH_OFFSET) as u32);
 /// Serialization errors.
 #[derive(Debug)]
 pub enum SerializeError {
-    /// The encoded length exceeds the maximum.
-    MaximumEncodedLengthExceeded(usize),
+    /// The encoded length was at least as big as the maximum length possible.
+    MaximumEncodedLengthReached(usize),
     /// An invalid instance was encountered.
     InvalidInstance(InstanceError),
     /// An invalid type was encountered.
@@ -34,9 +34,9 @@ impl From<TypeError> for SerializeError {
 impl Display for SerializeError {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
-            SerializeError::MaximumEncodedLengthExceeded(size) => write!(
+            SerializeError::MaximumEncodedLengthReached(size) => write!(
                 f,
-                "the encoded length is {size} which exceeds the maximum length {MAXIMUM_LENGTH}",
+                "the encoded length is {size} which meets or exceeds the maximum length {MAXIMUM_LENGTH}",
             ),
             SerializeError::InvalidInstance(err) => write!(f, "invalid instance: {err}"),
             SerializeError::InvalidType(err) => write!(f, "invalid type: {err}"),
@@ -55,71 +55,64 @@ pub trait Serialize {
     fn serialize(&self, buffer: &mut Vec<u8>) -> Result<usize, SerializeError>;
 }
 
-pub fn serialize_composite_from_components(
-    mut fixed: Vec<Option<Vec<u8>>>,
-    mut variable: Vec<u8>,
-    variable_lengths: Vec<usize>,
-    fixed_lengths_sum: usize,
-    buffer: &mut Vec<u8>,
-) -> Result<usize, SerializeError> {
-    debug_assert_eq!(fixed.len(), variable_lengths.len());
-
-    let total_size = fixed_lengths_sum + variable_lengths.iter().sum::<usize>();
-    if total_size as u64 >= MAXIMUM_LENGTH {
-        return Err(SerializeError::MaximumEncodedLengthExceeded(total_size))
-    }
+// Part represents either a fixed sized part of the serialization
+// or an offset pointing to a variably sized part of the serialization
+pub enum Part {
+    Fixed(Vec<u8>),
+    Offset(usize),
+}
 
-    // SAFETY: `fixed_lengths_sum` fits in `u32` if the total size check holds
-    let mut running_length = fixed_lengths_sum as u32;
-    debug_assert_eq!(fixed.len(), variable_lengths.len());
-    for (part_opt, variable_length) in fixed.iter_mut().zip(variable_lengths) {
-        if let Some(part) = part_opt {
-            buffer.append(part);
-        } else {
-            // SAFETY: `variable_length` fits in `u32` if the total size check holds
-            let bytes_written = running_length.serialize(buffer)?;
-            debug_assert_eq!(bytes_written, BYTES_PER_LENGTH_OFFSET);
+#[derive(Default)]
+pub struct Serializer {
+    parts: Vec<Part>,
+    variable: Vec<u8>,
+    fixed_lengths_sum: usize,
+    variable_lengths_sum: usize,
+}
 
-            running_length += variable_length as u32;
+impl Serializer {
+    pub fn serialize(mut self, buffer: &mut Vec<u8>) -> Result<usize, SerializeError> {
+        let total_size = self.fixed_lengths_sum + self.variable_lengths_sum;
+        if total_size as u64 >= MAXIMUM_LENGTH {
+            return Err(SerializeError::MaximumEncodedLengthReached(total_size))
         }
-    }
 
-    buffer.append(&mut variable);
+        // SAFETY: `fixed_lengths_sum` fits in `u32` if the total size check holds
+        let mut running_length = self.fixed_lengths_sum as u32;
+        for part in self.parts {
+            match part {
+                Part::Fixed(mut data) => {
+                    buffer.append(&mut data);
+                }
+                Part::Offset(offset) => {
+                    let bytes_written = running_length.serialize(buffer)?;
+                    debug_assert_eq!(bytes_written, BYTES_PER_LENGTH_OFFSET);
+
+                    // SAFETY: `offset` fits in `u32` if the total size check holds
+                    running_length += offset as u32;
+                }
+            }
+        }
 
-    Ok(total_size)
-}
+        buffer.append(&mut self.variable);
 
-pub fn serialize_composite<T: SimpleSerialize>(
-    elements: &[T],
-    buffer: &mut Vec<u8>,
-) -> Result<usize, SerializeError> {
-    let mut fixed = vec![];
-    let mut variable = vec![];
-    let mut variable_lengths = vec![];
-    let mut fixed_lengths_sum = 0;
+        Ok(total_size)
+    }
 
-    for element in elements {
+    pub fn with_element<T: SimpleSerialize>(&mut self, element: &T) -> Result<(), SerializeError> {
         let mut element_buffer = Vec::with_capacity(T::size_hint());
         element.serialize(&mut element_buffer)?;
 
         let element_buffer_len = element_buffer.len();
         if T::is_variable_size() {
-            fixed.push(None);
-            fixed_lengths_sum += BYTES_PER_LENGTH_OFFSET;
-            variable.append(&mut element_buffer);
-            variable_lengths.push(element_buffer_len);
+            self.parts.push(Part::Offset(element_buffer_len));
+            self.variable.append(&mut element_buffer);
+            self.fixed_lengths_sum += BYTES_PER_LENGTH_OFFSET;
+            self.variable_lengths_sum += element_buffer_len;
         } else {
-            fixed.push(Some(element_buffer));
-            fixed_lengths_sum += element_buffer_len;
-            variable_lengths.push(0)
+            self.parts.push(Part::Fixed(element_buffer));
+            self.fixed_lengths_sum += element_buffer_len;
         }
+        Ok(())
     }
-
-    serialize_composite_from_components(
-        fixed,
-        variable,
-        variable_lengths,
-        fixed_lengths_sum,
-        buffer,
-    )
 }

ssz-rs/src/vector.rs

@@ -3,7 +3,7 @@ use crate::{
     error::{Error, InstanceError, TypeError},
     lib::*,
     merkleization::{elements_to_chunks, merkleize, pack, MerkleizationError, Merkleized, Node},
-    ser::{serialize_composite, Serialize, SerializeError},
+    ser::{Serialize, SerializeError, Serializer},
     SimpleSerialize, Sized,
 };
 #[cfg(feature = "serde")]
@@ -181,7 +181,11 @@ where
         if N == 0 {
             return Err(TypeError::InvalidBound(N).into())
         }
-        serialize_composite(&self.data, buffer)
+        let mut serializer = Serializer::default();
+        for element in &self.data {
+            serializer.with_element(element)?;
+        }
+        serializer.serialize(buffer)
     }
 }