AK: Handle empty trailing chunks in `AllocatingMemoryStream::offset_of`

2 files changed, 37 insertions, 3 deletions

diff --git a/AK/MemoryStream.cpp b/AK/MemoryStream.cpp
index 674c7ec44c..c0d2be82c4 100644
--- a/AK/MemoryStream.cpp
+++ b/AK/MemoryStream.cpp
@@ -198,11 +198,12 @@ ErrorOr<Optional<size_t>> AllocatingMemoryStream::offset_of(ReadonlyBytes needle
     if (m_chunks.size() == 0)
         return Optional<size_t> {};
 
-    // Ensure that we don't have to trim away more than one block.
+    // Ensure that we don't have empty chunks at the beginning of the stream. Our trimming implementation
+    // assumes this to be the case, since this should be held up by `cleanup_unused_chunks()` at all times.
     VERIFY(m_read_offset < CHUNK_SIZE);
-    VERIFY(m_chunks.size() * CHUNK_SIZE - m_write_offset < CHUNK_SIZE);
 
-    auto chunk_count = m_chunks.size();
+    auto empty_chunks_at_end = ((m_chunks.size() * CHUNK_SIZE - m_write_offset) / CHUNK_SIZE);
+    auto chunk_count = m_chunks.size() - empty_chunks_at_end;
     auto search_spans = TRY(FixedArray<ReadonlyBytes>::create(chunk_count));
 
     for (size_t i = 0; i < chunk_count; i++) {
diff --git a/Tests/AK/TestMemoryStream.cpp b/Tests/AK/TestMemoryStream.cpp
index 097c92e4f2..073be8fd32 100644
--- a/Tests/AK/TestMemoryStream.cpp
+++ b/Tests/AK/TestMemoryStream.cpp
@@ -90,3 +90,36 @@ TEST_CASE(allocating_memory_stream_offset_of_oob)
         EXPECT(!offset.has_value());
     }
 }
+
+TEST_CASE(allocating_memory_stream_offset_of_after_chunk_reorder)
+{
+    AllocatingMemoryStream stream;
+
+    // First, fill exactly one chunk (in groups of 16 bytes). This chunk will be reordered.
+    for (size_t i = 0; i < AllocatingMemoryStream::CHUNK_SIZE / 16; ++i)
+        MUST(stream.write_until_depleted("AAAAAAAAAAAAAAAA"sv.bytes()));
+
+    // Append a few additional bytes to create a second chunk.
+    MUST(stream.write_until_depleted("BCDEFGHIJKLMNOPQ"sv.bytes()));
+
+    // Read back the first chunk, which should reorder it to the end of the list.
+    // The chunk that we wrote to the second time is now the first one.
+    MUST(stream.discard(AllocatingMemoryStream::CHUNK_SIZE));
+
+    {
+        auto offset = MUST(stream.offset_of("A"sv.bytes()));
+        EXPECT(!offset.has_value());
+    }
+
+    {
+        auto offset = MUST(stream.offset_of("B"sv.bytes()));
+        EXPECT(offset.has_value());
+        EXPECT_EQ(offset.value(), 0ul);
+    }
+
+    {
+        auto offset = MUST(stream.offset_of("Q"sv.bytes()));
+        EXPECT(offset.has_value());
+        EXPECT_EQ(offset.value(), 15ul);
+    }
+}