LibCore: Add an automatically allocating MemoryStream

Similar to AK::DuplexStream, this writes to the end and reads from the
beginning of a dynamically allocated buffer.

3 files changed, 200 insertions, 0 deletions

diff --git a/Tests/LibCore/TestLibCoreStream.cpp b/Tests/LibCore/TestLibCoreStream.cpp
index 40c7f92fd5..13d086efc4 100644
--- a/Tests/LibCore/TestLibCoreStream.cpp
+++ b/Tests/LibCore/TestLibCoreStream.cpp
@@ -5,8 +5,10 @@
  */
 
 #include <AK/Format.h>
+#include <AK/String.h>
 #include <LibCore/EventLoop.h>
 #include <LibCore/LocalServer.h>
+#include <LibCore/MemoryStream.h>
 #include <LibCore/Stream.h>
 #include <LibCore/TCPServer.h>
 #include <LibCore/Timer.h>
@@ -504,3 +506,54 @@ TEST_CASE(buffered_tcp_socket_read)
     auto second_received_line = maybe_second_received_line.value();
     EXPECT_EQ(second_received_line, second_line);
 }
+
+// Allocating memory stream tests
+
+TEST_CASE(allocating_memory_stream_empty)
+{
+    Core::Stream::AllocatingMemoryStream stream;
+
+    EXPECT_EQ(stream.used_buffer_size(), 0ul);
+
+    {
+        Array<u8, 32> array;
+        auto read_bytes = MUST(stream.read(array));
+        EXPECT_EQ(read_bytes.size(), 0ul);
+    }
+}
+
+TEST_CASE(allocating_memory_stream_10kb)
+{
+    auto file = MUST(Core::Stream::File::open("/usr/Tests/LibCore/10kb.txt"sv, Core::Stream::OpenMode::Read));
+    size_t const file_size = MUST(file->size());
+    size_t constexpr test_chunk_size = 4096;
+
+    // Read file contents into the memory stream.
+    Core::Stream::AllocatingMemoryStream stream;
+    while (!file->is_eof()) {
+        Array<u8, test_chunk_size> array;
+        MUST(stream.write(MUST(file->read(array))));
+    }
+
+    EXPECT_EQ(stream.used_buffer_size(), file_size);
+
+    MUST(file->seek(0, Core::Stream::SeekMode::SetPosition));
+
+    // Check the stream contents when reading back.
+    size_t offset = 0;
+    while (!file->is_eof()) {
+        Array<u8, test_chunk_size> file_array;
+        Array<u8, test_chunk_size> stream_array;
+        auto file_span = MUST(file->read(file_array));
+        auto stream_span = MUST(stream.read(stream_array));
+        EXPECT_EQ(file_span.size(), stream_span.size());
+
+        for (size_t i = 0; i < file_span.size(); i++) {
+            if (file_array[i] == stream_array[i])
+                continue;
+
+            FAIL(String::formatted("Data started to diverge at index {}: file={}, stream={}", offset + i, file_array[i], stream_array[i]));
+        }
+        offset += file_span.size();
+    }
+}
diff --git a/Userland/Libraries/LibCore/MemoryStream.cpp b/Userland/Libraries/LibCore/MemoryStream.cpp
index 0a9d0d6f07..b5ed036ce7 100644
--- a/Userland/Libraries/LibCore/MemoryStream.cpp
+++ b/Userland/Libraries/LibCore/MemoryStream.cpp
@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, kleines Filmröllchen <filmroellchen@serenityos.org>.
+ * Copyright (c) 2022, Tim Schumacher <timschumi@gmx.de>.
  *
  * SPDX-License-Identifier: BSD-2-Clause
  */
@@ -124,4 +125,123 @@ size_t FixedMemoryStream::remaining() const
     return m_bytes.size() - m_offset;
 }
 
+ErrorOr<Bytes> AllocatingMemoryStream::read(Bytes bytes)
+{
+    size_t read_bytes = 0;
+
+    while (read_bytes < bytes.size()) {
+        VERIFY(m_write_offset >= m_read_offset);
+
+        auto range = TRY(next_read_range());
+        if (range.size() == 0)
+            break;
+
+        auto copied_bytes = range.copy_trimmed_to(bytes.slice(read_bytes));
+
+        read_bytes += copied_bytes;
+        m_read_offset += copied_bytes;
+    }
+
+    cleanup_unused_chunks();
+
+    return bytes.trim(read_bytes);
+}
+
+ErrorOr<size_t> AllocatingMemoryStream::write(ReadonlyBytes bytes)
+{
+    size_t written_bytes = 0;
+
+    while (written_bytes < bytes.size()) {
+        VERIFY(m_write_offset >= m_read_offset);
+
+        auto range = TRY(next_write_range());
+
+        auto copied_bytes = bytes.slice(written_bytes).copy_trimmed_to(range);
+
+        written_bytes += copied_bytes;
+        m_write_offset += copied_bytes;
+    }
+
+    return written_bytes;
+}
+
+ErrorOr<void> AllocatingMemoryStream::discard(size_t count)
+{
+    VERIFY(m_write_offset >= m_read_offset);
+
+    if (count > used_buffer_size())
+        return Error::from_string_literal("Number of discarded bytes is higher than the number of allocated bytes");
+
+    m_read_offset += count;
+
+    cleanup_unused_chunks();
+
+    return {};
+}
+
+bool AllocatingMemoryStream::is_eof() const
+{
+    return used_buffer_size() == 0;
+}
+
+bool AllocatingMemoryStream::is_open() const
+{
+    return true;
+}
+
+void AllocatingMemoryStream::close()
+{
+}
+
+size_t AllocatingMemoryStream::used_buffer_size() const
+{
+    return m_write_offset - m_read_offset;
+}
+
+ErrorOr<ReadonlyBytes> AllocatingMemoryStream::next_read_range()
+{
+    VERIFY(m_write_offset >= m_read_offset);
+
+    size_t const chunk_index = m_read_offset / chunk_size;
+    size_t const chunk_offset = m_read_offset % chunk_size;
+    size_t const read_size = min(chunk_size - m_read_offset % chunk_size, m_write_offset - m_read_offset);
+
+    if (read_size == 0)
+        return ReadonlyBytes { static_cast<u8*>(nullptr), 0 };
+
+    VERIFY(chunk_index < m_chunks.size());
+
+    return ReadonlyBytes { m_chunks[chunk_index].data() + chunk_offset, read_size };
+}
+
+ErrorOr<Bytes> AllocatingMemoryStream::next_write_range()
+{
+    VERIFY(m_write_offset >= m_read_offset);
+
+    size_t const chunk_index = m_write_offset / chunk_size;
+    size_t const chunk_offset = m_write_offset % chunk_size;
+    size_t const write_size = chunk_size - m_write_offset % chunk_size;
+
+    if (chunk_index >= m_chunks.size())
+        TRY(m_chunks.try_append(TRY(Chunk::create_uninitialized(chunk_size))));
+
+    VERIFY(chunk_index < m_chunks.size());
+
+    return Bytes { m_chunks[chunk_index].data() + chunk_offset, write_size };
+}
+
+void AllocatingMemoryStream::cleanup_unused_chunks()
+{
+    // FIXME: Move these all at once.
+    while (m_read_offset >= chunk_size) {
+        VERIFY(m_write_offset >= m_read_offset);
+
+        auto buffer = m_chunks.take_first();
+        m_read_offset -= chunk_size;
+        m_write_offset -= chunk_size;
+
+        m_chunks.append(move(buffer));
+    }
+}
+
 }
diff --git a/Userland/Libraries/LibCore/MemoryStream.h b/Userland/Libraries/LibCore/MemoryStream.h
index ee6b6e19ee..902c19e2b5 100644
--- a/Userland/Libraries/LibCore/MemoryStream.h
+++ b/Userland/Libraries/LibCore/MemoryStream.h
@@ -47,4 +47,31 @@ private:
     bool m_writing_enabled { true };
 };
 
+/// A stream class that allows for writing to an automatically allocating memory area
+/// and reading back the written data afterwards.
+class AllocatingMemoryStream final : public Stream {
+public:
+    virtual ErrorOr<Bytes> read(Bytes) override;
+    virtual ErrorOr<size_t> write(ReadonlyBytes) override;
+    virtual ErrorOr<void> discard(size_t) override;
+    virtual bool is_eof() const override;
+    virtual bool is_open() const override;
+    virtual void close() override;
+
+    size_t used_buffer_size() const;
+
+private:
+    // Note: We set the inline buffer capacity to zero to make moving chunks as efficient as possible.
+    using Chunk = AK::Detail::ByteBuffer<0>;
+    static constexpr size_t chunk_size = 4096;
+
+    ErrorOr<ReadonlyBytes> next_read_range();
+    ErrorOr<Bytes> next_write_range();
+    void cleanup_unused_chunks();
+
+    Vector<Chunk> m_chunks;
+    size_t m_read_offset = 0;
+    size_t m_write_offset = 0;
+};
+
 }