Merge pull request #5 from MaterializeInc/fallible-allocations

decoding: Make allocations fallible

src/encoding.rs

@@ -4,7 +4,7 @@
 
 #![allow(clippy::implicit_hasher, clippy::ptr_arg)]
 
-use alloc::collections::BTreeMap;
+use alloc::collections::{BTreeMap, TryReserveError};
 use alloc::format;
 use alloc::string::String;
 use alloc::vec::Vec;
@@ -452,6 +452,7 @@ macro_rules! merge_repeated_numeric {
                 merge_loop(values, buf, ctx, |values, buf, ctx| {
                     let mut value = Default::default();
                     $merge($wire_type, &mut value, buf, ctx)?;
+                    values.try_reserve(1)?;
                     values.push(value);
                     Ok(())
                 })
@@ -460,6 +461,7 @@ macro_rules! merge_repeated_numeric {
                 check_wire_type($wire_type, wire_type)?;
                 let mut value = Default::default();
                 $merge(wire_type, &mut value, buf, ctx)?;
+                values.try_reserve(1)?;
                 values.push(value);
                 Ok(())
             }
@@ -771,7 +773,15 @@ macro_rules! length_delimited {
             check_wire_type(WireType::LengthDelimited, wire_type)?;
             let mut value = Default::default();
             merge(wire_type, &mut value, buf, ctx)?;
+
+            // Explicitly reserve before pushing so we can return an error instead of panicking from
+            // and out-of-memory exception.
+            //
+            // Note: Reserving a single element will still cause ammortizied growth of the Vec like
+            // pushing does, so there is no performance impact of reserving early.
+            values.try_reserve(1)?;
             values.push(value);
+
             Ok(())
         }
 
@@ -877,12 +887,13 @@ pub trait BytesAdapter: sealed::BytesAdapter {}
 
 mod sealed {
     use super::{Buf, BufMut};
+    use alloc::collections::TryReserveError;
 
     pub trait BytesAdapter: Default + Sized + 'static {
         fn len(&self) -> usize;
 
         /// Replace contents of this buffer with the contents of another buffer.
-        fn replace_with<B>(&mut self, buf: B)
+        fn replace_with<B>(&mut self, buf: B) -> Result<(), TryReserveError>
         where
             B: Buf;
 
@@ -904,11 +915,12 @@ impl sealed::BytesAdapter for Bytes {
         Buf::remaining(self)
     }
 
-    fn replace_with<B>(&mut self, mut buf: B)
+    fn replace_with<B>(&mut self, mut buf: B) -> Result<(), TryReserveError>
     where
         B: Buf,
     {
         *self = buf.copy_to_bytes(buf.remaining());
+        Ok(())
     }
 
     fn append_to<B>(&self, buf: &mut B)
@@ -926,13 +938,14 @@ impl sealed::BytesAdapter for Vec<u8> {
         Vec::len(self)
     }
 
-    fn replace_with<B>(&mut self, buf: B)
+    fn replace_with<B>(&mut self, buf: B) -> Result<(), TryReserveError>
     where
         B: Buf,
     {
         self.clear();
-        self.reserve(buf.remaining());
+        self.try_reserve(buf.remaining())?;
         self.put(buf);
+        Ok(())
     }
 
     fn append_to<B>(&self, buf: &mut B)
@@ -985,11 +998,11 @@ pub mod bytes {
         // This is intended for A and B both being Bytes so it is zero-copy.
         // Some combinations of A and B types may cause a double-copy,
         // in which case merge_one_copy() should be used instead.
-        value.replace_with(buf.copy_to_bytes(len));
+        value.replace_with(buf.copy_to_bytes(len))?;
         Ok(())
     }
 
-    pub(super) fn merge_one_copy<A, B>(
+    pub fn merge_one_copy<A, B>(
         wire_type: WireType,
         value: &mut A,
         buf: &mut B,
@@ -1007,7 +1020,7 @@ pub mod bytes {
         let len = len as usize;
 
         // If we must copy, make sure to copy only once.
-        value.replace_with(buf.take(len));
+        value.replace_with(buf.take(len))?;
         Ok(())
     }
 
@@ -1111,7 +1124,15 @@ pub mod message {
         check_wire_type(WireType::LengthDelimited, wire_type)?;
         let mut msg = M::default();
         merge(WireType::LengthDelimited, &mut msg, buf, ctx)?;
+
+        // Explicitly reserve before pushing so we can return an error instead of panicking from
+        // and out-of-memory exception.
+        //
+        // Note: Reserving a single element will still cause ammortizied growth of the Vec like
+        // pushing does, so there is no performance impact of reserving early.
+        messages.try_reserve(1)?;
         messages.push(msg);
+
         Ok(())
     }
 
@@ -1202,7 +1223,15 @@ pub mod group {
         check_wire_type(WireType::StartGroup, wire_type)?;
         let mut msg = M::default();
         merge(tag, WireType::StartGroup, &mut msg, buf, ctx)?;
+
+        // Explicitly reserve before pushing so we can return an error instead of panicking from
+        // and out-of-memory exception.
+        //
+        // Note: Reserving a single element will still cause ammortizied growth of the Vec like
+        // pushing does, so there is no performance impact of reserving early.
+        messages.try_reserve(1)?;
         messages.push(msg);
+
         Ok(())
     }
 

src/error.rs

@@ -2,6 +2,7 @@
 
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
+use alloc::collections::TryReserveError;
 use alloc::vec::Vec;
 
 use core::fmt;
@@ -69,6 +70,12 @@ impl fmt::Display for DecodeError {
     }
 }
 
+impl From<TryReserveError> for DecodeError {
+    fn from(error: TryReserveError) -> Self {
+        DecodeError::new(error.to_string())
+    }
+}
+
 #[cfg(feature = "std")]
 impl std::error::Error for DecodeError {}
 

src/types.rs

@@ -343,7 +343,7 @@ impl Message for Vec<u8> {
         B: Buf,
     {
         if tag == 1 {
-            bytes::merge(wire_type, self, buf, ctx)
+            bytes::merge_one_copy(wire_type, self, buf, ctx)
         } else {
             skip_field(wire_type, tag, buf, ctx)
         }