From f2731c0788ab2a831c8cd795c3f0eee15dfc631b Mon Sep 17 00:00:00 2001
From: Victor Koenders <gpg@trangar.com>
Date: Sat, 11 Dec 2021 11:00:04 +0100
Subject: [PATCH] Added try_reserve to VecWriter, Vec, VecDeque and HashMap

---
 src/error.rs               | 17 +++++++++++++++++
 src/features/impl_alloc.rs | 17 ++++++++++++++---
 src/features/impl_std.rs   |  5 ++++-
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index b9c6a3b6..aac35007 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,5 +1,8 @@
 //! Errors that can be encounting by Encoding and Decoding.
 
+#[cfg(feature = "alloc")]
+use alloc::collections::TryReserveError;
+
 /// Errors that can be encountered by encoding a type
 #[non_exhaustive]
 #[derive(Debug)]
@@ -54,6 +57,13 @@ pub enum EncodeError {
     /// Serde provided bincode with a sequence without a length, which is not supported in bincode
     #[cfg(feature = "serde")]
     SequenceMustHaveLength,
+
+    /// bincode failed to allocate enough memory
+    #[cfg(feature = "alloc")]
+    OutOfMemory {
+        /// The inner error
+        inner: TryReserveError,
+    },
 }
 
 impl core::fmt::Display for EncodeError {
@@ -139,6 +149,13 @@ pub enum DecodeError {
     /// Serde tried decoding a borrowed value from an owned reader. Use `serde_decode_borrowed_from_*` instead
     #[cfg(feature = "serde")]
     CannotBorrowOwnedData,
+
+    /// bincode failed to allocate enough memory
+    #[cfg(feature = "alloc")]
+    OutOfMemory {
+        /// The inner error
+        inner: TryReserveError,
+    },
 }
 
 impl core::fmt::Display for DecodeError {
diff --git a/src/features/impl_alloc.rs b/src/features/impl_alloc.rs
index 366aa735..b05487fa 100644
--- a/src/features/impl_alloc.rs
+++ b/src/features/impl_alloc.rs
@@ -30,6 +30,9 @@ impl VecWriter {
 
 impl enc::write::Writer for VecWriter {
     fn write(&mut self, bytes: &[u8]) -> Result<(), EncodeError> {
+        self.inner
+            .try_reserve(bytes.len())
+            .map_err(|inner| EncodeError::OutOfMemory { inner })?;
         self.inner.extend_from_slice(bytes);
         Ok(())
     }
@@ -52,7 +55,11 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        let mut map = BinaryHeap::with_capacity(len);
+        let mut map = BinaryHeap::new();
+        // TODO:
+        // map.try_reserve(len).map_err(|inner| DecodeError::OutOfMemory { inner })?;
+        map.reserve(len);
+
         for _ in 0..len {
             let key = T::decode(&mut decoder)?;
             map.push(key);
@@ -140,7 +147,9 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        let mut map = VecDeque::with_capacity(len);
+        let mut map = VecDeque::new();
+        map.try_reserve(len)
+            .map_err(|inner| DecodeError::OutOfMemory { inner })?;
         for _ in 0..len {
             let key = T::decode(&mut decoder)?;
             map.push_back(key);
@@ -168,7 +177,9 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        let mut vec = Vec::with_capacity(len);
+        let mut vec = Vec::new();
+        vec.try_reserve(len)
+            .map_err(|inner| DecodeError::OutOfMemory { inner })?;
         for _ in 0..len {
             vec.push(T::decode(&mut decoder)?);
         }
diff --git a/src/features/impl_std.rs b/src/features/impl_std.rs
index fe53cdac..f6eee5d9 100644
--- a/src/features/impl_std.rs
+++ b/src/features/impl_std.rs
@@ -369,7 +369,10 @@ where
 {
     fn decode<D: Decoder>(mut decoder: D) -> Result<Self, DecodeError> {
         let len = crate::de::decode_slice_len(&mut decoder)?;
-        let mut map = HashMap::with_capacity(len);
+        let mut map = HashMap::new();
+        map.try_reserve(len)
+            .map_err(|inner| DecodeError::OutOfMemory { inner })?;
+
         for _ in 0..len {
             let k = K::decode(&mut decoder)?;
             let v = V::decode(&mut decoder)?;