From c1e6e67881c5b3d3260ee83bd403bc6bc78f3edb 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 17fee056..b1c10387 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)]
@@ -51,6 +54,13 @@ pub enum EncodeError {
         time: std::time::SystemTime,
     },
 
+    /// bincode failed to allocate enough memory
+    #[cfg(feature = "alloc")]
+    OutOfMemory {
+        /// The inner error
+        inner: TryReserveError,
+    },
+
     #[cfg(feature = "serde")]
     /// A serde-specific error that occurred while decoding.
     Serde(crate::features::serde::EncodeError),
@@ -168,6 +178,13 @@ pub enum DecodeError {
     #[cfg(feature = "serde")]
     /// A serde-specific error that occurred while decoding.
     Serde(crate::features::serde::DecodeError),
+
+    /// 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 ea97cd6d..894cde93 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(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(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(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 811148ef..5d690684 100644
--- a/src/features/impl_std.rs
+++ b/src/features/impl_std.rs
@@ -383,7 +383,10 @@ where
         let len = crate::de::decode_slice_len(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)>());