From 247f7c5fcc1f2384940735ef7939f428cabf79d4 Mon Sep 17 00:00:00 2001
From: Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
Date: Mon, 14 Apr 2025 03:17:21 +0200
Subject: [PATCH] AK: Add peek_some() to AllocatingMemoryStream

Same as read_some(), but doesn't move read position.
---
 AK/MemoryStream.cpp | 28 ++++++++++++++++++++++------
 AK/MemoryStream.h   |  4 +++-
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/AK/MemoryStream.cpp b/AK/MemoryStream.cpp
index a4baeb819e6..bbc7b19947d 100644
--- a/AK/MemoryStream.cpp
+++ b/AK/MemoryStream.cpp
@@ -124,6 +124,22 @@ size_t FixedMemoryStream::remaining() const
     return m_bytes.size() - m_offset;
 }
 
+void AllocatingMemoryStream::peek_some(Bytes bytes) const
+{
+    size_t read_bytes = 0;
+
+    auto peek_offset = m_read_offset;
+    while (read_bytes < bytes.size()) {
+        auto range = MUST(next_read_range(peek_offset));
+        if (range.size() == 0)
+            break;
+
+        auto copied_bytes = range.copy_trimmed_to(bytes.slice(read_bytes));
+        read_bytes += copied_bytes;
+        peek_offset += copied_bytes;
+    }
+}
+
 ErrorOr<Bytes> AllocatingMemoryStream::read_some(Bytes bytes)
 {
     size_t read_bytes = 0;
@@ -131,7 +147,7 @@ ErrorOr<Bytes> AllocatingMemoryStream::read_some(Bytes bytes)
     while (read_bytes < bytes.size()) {
         VERIFY(m_write_offset >= m_read_offset);
 
-        auto range = TRY(next_read_range());
+        auto range = TRY(next_read_range(m_read_offset));
         if (range.size() == 0)
             break;
 
@@ -231,13 +247,13 @@ ErrorOr<Optional<size_t>> AllocatingMemoryStream::offset_of(ReadonlyBytes needle
     return AK::memmem(search_spans.begin(), search_spans.end(), needle);
 }
 
-ErrorOr<ReadonlyBytes> AllocatingMemoryStream::next_read_range()
+ErrorOr<ReadonlyBytes> AllocatingMemoryStream::next_read_range(size_t read_offset) const
 {
-    VERIFY(m_write_offset >= m_read_offset);
+    VERIFY(m_write_offset >= 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);
+    size_t const chunk_index = read_offset / CHUNK_SIZE;
+    size_t const chunk_offset = read_offset % CHUNK_SIZE;
+    size_t const read_size = min(CHUNK_SIZE - read_offset % CHUNK_SIZE, m_write_offset - read_offset);
 
     if (read_size == 0)
         return ReadonlyBytes { static_cast<u8*>(nullptr), 0 };
diff --git a/AK/MemoryStream.h b/AK/MemoryStream.h
index 192de0febbe..d5ca9c8a114 100644
--- a/AK/MemoryStream.h
+++ b/AK/MemoryStream.h
@@ -86,6 +86,8 @@ class AllocatingMemoryStream final : public Stream {
 public:
     static constexpr size_t CHUNK_SIZE = 4096;
 
+    void peek_some(Bytes) const;
+
     virtual ErrorOr<Bytes> read_some(Bytes) override;
     virtual ErrorOr<size_t> write_some(ReadonlyBytes) override;
     virtual ErrorOr<void> discard(size_t) override;
@@ -101,7 +103,7 @@ private:
     // Note: We set the inline buffer capacity to zero to make moving chunks as efficient as possible.
     using Chunk = AK::Detail::ByteBuffer<0>;
 
-    ErrorOr<ReadonlyBytes> next_read_range();
+    ErrorOr<ReadonlyBytes> next_read_range(size_t read_offset) const;
     ErrorOr<Bytes> next_write_range();
     void cleanup_unused_chunks();