From 8466fff1cbea429602492371b5fa79588b1a792d 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 | 19 ++++++++++++++++---
 src/features/impl_std.rs   |  5 ++++-
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 70465c03..8a1b7aba 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 {
@@ -148,6 +158,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 e8316f1a..1d354039 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(())
     }
@@ -54,7 +57,11 @@ where
         let len = crate::de::decode_slice_len(&mut decoder)?;
         decoder.claim_container_read::<T>(len)?;
 
-        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 {
             // See the documentation on `unclaim_bytes_read` as to why we're doing this here
             decoder.unclaim_bytes_read(core::mem::size_of::<T>());
@@ -157,7 +164,10 @@ where
         let len = crate::de::decode_slice_len(&mut decoder)?;
         decoder.claim_container_read::<T>(len)?;
 
-        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 {
             // See the documentation on `unclaim_bytes_read` as to why we're doing this here
             decoder.unclaim_bytes_read(core::mem::size_of::<T>());
@@ -190,7 +200,10 @@ where
         let len = crate::de::decode_slice_len(&mut decoder)?;
         decoder.claim_container_read::<T>(len)?;
 
-        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 {
             // See the documentation on `unclaim_bytes_read` as to why we're doing this here
             decoder.unclaim_bytes_read(core::mem::size_of::<T>());
diff --git a/src/features/impl_std.rs b/src/features/impl_std.rs
index 9f4e1039..4c06635e 100644
--- a/src/features/impl_std.rs
+++ b/src/features/impl_std.rs
@@ -371,7 +371,10 @@ where
         let len = crate::de::decode_slice_len(&mut decoder)?;
         decoder.claim_container_read::<(K, V)>(len)?;
 
-        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 {
             // See the documentation on `unclaim_bytes_read` as to why we're doing this here
             decoder.unclaim_bytes_read(core::mem::size_of::<(K, V)>());