AK: Re-add OutputMemoryStream for static buffers only.

3 files changed, 66 insertions, 4 deletions

diff --git a/AK/Forward.h b/AK/Forward.h
index c950a29063..dc73be66bd 100644
--- a/AK/Forward.h
+++ b/AK/Forward.h
@@ -53,6 +53,7 @@ class InputMemoryStream;
 class DuplexMemoryStream;
 class OutputStream;
 class InputBitStream;
+class OutputMemoryStream;
 
 template<size_t Capacity>
 class CircularDuplexStream;
@@ -153,6 +154,7 @@ using AK::LogStream;
 using AK::NonnullOwnPtr;
 using AK::NonnullRefPtr;
 using AK::Optional;
+using AK::OutputMemoryStream;
 using AK::OutputStream;
 using AK::OwnPtr;
 using AK::ReadonlyBytes;
diff --git a/AK/MemoryStream.h b/AK/MemoryStream.h
index 0b1a2d311a..103a72b8f1 100644
--- a/AK/MemoryStream.h
+++ b/AK/MemoryStream.h
@@ -35,7 +35,7 @@ namespace AK {
 
 class InputMemoryStream final : public InputStream {
 public:
-    InputMemoryStream(ReadonlyBytes bytes)
+    explicit InputMemoryStream(ReadonlyBytes bytes)
         : m_bytes(bytes)
     {
     }
@@ -161,9 +161,45 @@ private:
     size_t m_offset { 0 };
 };
 
-// All data written to this stream can be read from it. Reading and writing is done
-// using different offsets, meaning that it is not necessary to seek to the start
-// before reading; this behaviour differs from BufferStream.
+class OutputMemoryStream final : public OutputStream {
+public:
+    explicit OutputMemoryStream(Bytes bytes)
+        : m_bytes(bytes)
+    {
+    }
+
+    size_t write(ReadonlyBytes bytes) override
+    {
+        const auto nwritten = bytes.copy_trimmed_to(m_bytes.slice(m_offset));
+        m_offset += nwritten;
+        return nwritten;
+    }
+
+    bool write_or_error(ReadonlyBytes bytes) override
+    {
+        if (remaining() < bytes.size()) {
+            set_recoverable_error();
+            return false;
+        }
+
+        write(bytes);
+        return true;
+    }
+
+    ReadonlyBytes bytes() const { return { data(), size() }; }
+    Bytes bytes() { return { data(), size() }; }
+
+    const u8* data() const { return m_bytes.data(); }
+    u8* data() { return m_bytes.data(); }
+
+    size_t size() const { return m_offset; }
+    size_t remaining() const { return m_bytes.size() - m_offset; }
+
+private:
+    size_t m_offset { 0 };
+    Bytes m_bytes;
+};
+
 class DuplexMemoryStream final : public DuplexStream {
 public:
     static constexpr size_t chunk_size = 4 * 1024;
@@ -321,3 +357,4 @@ private:
 using AK::DuplexMemoryStream;
 using AK::InputMemoryStream;
 using AK::InputStream;
+using AK::OutputMemoryStream;
diff --git a/AK/Tests/TestMemoryStream.cpp b/AK/Tests/TestMemoryStream.cpp
index da419856be..789e9c488e 100644
--- a/AK/Tests/TestMemoryStream.cpp
+++ b/AK/Tests/TestMemoryStream.cpp
@@ -171,4 +171,27 @@ TEST_CASE(write_endian_values)
     EXPECT(compare({ expected, sizeof(expected) }, stream.copy_into_contiguous_buffer()));
 }
 
+TEST_CASE(new_output_memory_stream)
+{
+    Array<u8, 16> buffer;
+    OutputMemoryStream stream { buffer };
+
+    EXPECT_EQ(stream.size(), 0u);
+    EXPECT_EQ(stream.remaining(), 16u);
+
+    stream << LittleEndian<u16>(0x12'87);
+
+    EXPECT_EQ(stream.size(), 2u);
+    EXPECT_EQ(stream.remaining(), 14u);
+
+    stream << buffer;
+
+    EXPECT(stream.handle_recoverable_error());
+    EXPECT_EQ(stream.size(), 2u);
+    EXPECT_EQ(stream.remaining(), 14u);
+
+    EXPECT_EQ(stream.bytes().data(), buffer.data());
+    EXPECT_EQ(stream.bytes().size(), 2u);
+}
+
 TEST_MAIN(MemoryStream)