Allow byte_stream->byte_slice conversion to shrink unused buffer space

3 files changed, 24 insertions, 6 deletions

diff --git a/contrib/epee/include/byte_slice.h b/contrib/epee/include/byte_slice.h
index 6b79f6d92..18d60e088 100644
--- a/contrib/epee/include/byte_slice.h
+++ b/contrib/epee/include/byte_slice.h
@@ -112,7 +112,7 @@ namespace epee
     explicit byte_slice(std::string&& buffer);
 
     //! Convert `stream` into a slice with zero allocations.
-    explicit byte_slice(byte_stream&& stream) noexcept;
+    explicit byte_slice(byte_stream&& stream, bool shrink = true);
 
     byte_slice(byte_slice&& source) noexcept;
     ~byte_slice() noexcept = default;
diff --git a/contrib/epee/src/byte_slice.cpp b/contrib/epee/src/byte_slice.cpp
index faf7689be..453b63a4c 100644
--- a/contrib/epee/src/byte_slice.cpp
+++ b/contrib/epee/src/byte_slice.cpp
@@ -36,6 +36,11 @@
 #include "byte_slice.h"
 #include "byte_stream.h"
 
+namespace
+{
+  const std::size_t page_size = 4096;
+}
+
 namespace epee
 {
   struct byte_slice_data
@@ -173,16 +178,27 @@ namespace epee
     : byte_slice(adapt_buffer{}, std::move(buffer))
   {}
 
-  byte_slice::byte_slice(byte_stream&& stream) noexcept
+  byte_slice::byte_slice(byte_stream&& stream, const bool shrink)
     : storage_(nullptr), portion_(stream.data(), stream.size())
   {
-    if (stream.size())
+    if (portion_.size())
     {
-      std::uint8_t* const data = stream.take_buffer().release() - sizeof(raw_byte_slice);
+      byte_buffer buf;
+      if (shrink && page_size <= stream.available())
+      {
+          buf = byte_buffer_resize(stream.take_buffer(), portion_.size());
+          if (!buf)
+            throw std::bad_alloc{};
+          portion_ = {buf.get(), portion_.size()};
+      }
+      else // no need to shrink buffer
+        buf = stream.take_buffer();
+
+      std::uint8_t* const data = buf.release() - sizeof(raw_byte_slice);
       new (data) raw_byte_slice{};
       storage_.reset(reinterpret_cast<raw_byte_slice*>(data));
     }
-    else
+    else // empty stream
       portion_ = nullptr;
   }
 
diff --git a/tests/unit_tests/epee_utils.cpp b/tests/unit_tests/epee_utils.cpp
index b365cad86..0a1f0d3de 100644
--- a/tests/unit_tests/epee_utils.cpp
+++ b/tests/unit_tests/epee_utils.cpp
@@ -1090,11 +1090,13 @@ TEST(ByteStream, ToByteSlice)
 
   epee::byte_stream stream;
 
+  stream.reserve(128*1024);
   stream.write(source);
   EXPECT_EQ(sizeof(source), stream.size());
+  EXPECT_EQ(128*1024, stream.capacity());
   EXPECT_TRUE(equal(source, byte_span{stream.data(), stream.size()}));
 
-  const epee::byte_slice slice{std::move(stream)};
+  const epee::byte_slice slice{std::move(stream), true};
   EXPECT_EQ(0u, stream.size());
   EXPECT_EQ(0u, stream.available());
   EXPECT_EQ(0u, stream.capacity());