From 28f8a39726f6c94c28f06147a2846274a416d5a7 Mon Sep 17 00:00:00 2001
From: Gabe Black <gabe.black@gmail.com>
Date: Thu, 24 Dec 2020 06:05:14 -0800
Subject: [PATCH] base: Add a setNext method to the ChunkGenerator.

This method lets you stretch the current chunk, if you want to skip over
some of the upcoming chunks since you've already handled their bytes.

For instance, if you were iterating over pages in a range of virtual
addresses, you might be able to handle multiple page sized chunks at
once if they were represented by a single large sized page table entry.
This mechanism would let you move past all the pages you had just
handled without having to walk through them all one by one.

Change-Id: I7d962f548947b77f0aa1b725036dbcc9e1b28659
Reviewed-on: https://gem5-review.googlesource.com/c/public/gem5/+/38718
Reviewed-by: Gabe Black <gabe.black@gmail.com>
Maintainer: Gabe Black <gabe.black@gmail.com>
Tested-by: kokoro <noreply+kokoro@google.com>
---
 src/base/chunk_generator.hh      | 37 ++++++++++++++++++++++++++++---
 src/base/chunk_generator.test.cc | 38 ++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/src/base/chunk_generator.hh b/src/base/chunk_generator.hh
index 4c749b296..e95f4fda0 100644
--- a/src/base/chunk_generator.hh
+++ b/src/base/chunk_generator.hh
@@ -35,6 +35,7 @@
  */
 
 #include <algorithm>
+#include <cassert>
 
 #include "base/intmath.hh"
 #include "base/types.hh"
@@ -61,6 +62,8 @@ class ChunkGenerator
     Addr nextAddr;
     /** The size of the current chunk (in bytes). */
     Addr curSize;
+    /** The size of the next chunk (in bytes). */
+    Addr nextSize;
     /** The number of bytes remaining in the region after the current chunk. */
     Addr sizeLeft;
     /** The start address so we can calculate offset in writing block. */
@@ -96,12 +99,13 @@ class ChunkGenerator
                 // ... even if startAddr is already chunk-aligned
                 nextAddr += chunkSize;
             }
+            nextAddr = std::min(nextAddr, startAddr + totalSize);
         }
 
         // How many bytes are left between curAddr and the end of this chunk?
-        Addr left_in_chunk = nextAddr - curAddr;
-        curSize = std::min(totalSize, left_in_chunk);
+        curSize = nextAddr - curAddr;
         sizeLeft = totalSize - curSize;
+        nextSize = std::min(sizeLeft, chunkSize);
     }
 
     /**
@@ -141,6 +145,32 @@ class ChunkGenerator
      */
     bool last() const { return sizeLeft == 0; }
 
+    /**
+     * Grow this chunk to cover additional bytes which are already handled.
+     * @param next The first byte of the next chunk.
+     *
+     * @ingroup api_chunk_generator
+     */
+    void
+    setNext(Addr next)
+    {
+        assert(next >= nextAddr);
+
+        const Addr skipping = std::min(next - nextAddr, sizeLeft);
+
+        sizeLeft -= skipping;
+        curSize += skipping;
+        nextAddr = next;
+
+        assert(chunkSize);
+
+        // nextSize will be enough to get to an alignment boundary,
+        nextSize = roundUp(next, chunkSize) - next;
+        // or if it's already aligned, to the following boundary or the end.
+        if (!nextSize)
+            nextSize = std::min(sizeLeft, chunkSize);
+    }
+
     /**
      * Advance generator to next chunk.
      * @return True if successful, false if unsuccessful
@@ -157,9 +187,10 @@ class ChunkGenerator
         }
 
         curAddr = nextAddr;
-        curSize = std::min(sizeLeft, chunkSize);
+        curSize = nextSize;
         sizeLeft -= curSize;
         nextAddr += curSize;
+        nextSize = std::min(sizeLeft, chunkSize);
         return true;
     }
 };
diff --git a/src/base/chunk_generator.test.cc b/src/base/chunk_generator.test.cc
index 91895cbc3..f5d7048b0 100644
--- a/src/base/chunk_generator.test.cc
+++ b/src/base/chunk_generator.test.cc
@@ -59,6 +59,44 @@ TEST(ChunkGeneratorTest, AdvanceToNextChunk)
     EXPECT_FALSE(chunk_generator.last());
 }
 
+/*
+ * A test to check skipping over bytes.
+ */
+TEST(ChunkGeneratorTest, SkipBytes)
+{
+    ChunkGenerator chunk_generator(0, 1024, 8);
+    EXPECT_EQ(0, chunk_generator.addr());
+    EXPECT_TRUE(chunk_generator.next());
+    EXPECT_EQ(8, chunk_generator.addr());
+
+    chunk_generator.setNext(23);
+    EXPECT_EQ(23 - 8, chunk_generator.size());
+    EXPECT_TRUE(chunk_generator.next());
+    EXPECT_EQ(23, chunk_generator.addr());
+    EXPECT_EQ(1, chunk_generator.size());
+    EXPECT_TRUE(chunk_generator.next());
+    EXPECT_EQ(24, chunk_generator.addr());
+    EXPECT_EQ(8, chunk_generator.size());
+
+    chunk_generator.setNext(32);
+    EXPECT_EQ(32 - 24, chunk_generator.size());
+    EXPECT_TRUE(chunk_generator.next());
+    EXPECT_EQ(32, chunk_generator.addr());
+    EXPECT_EQ(8, chunk_generator.size());
+
+    chunk_generator.setNext(64);
+    EXPECT_EQ(64 - 32, chunk_generator.size());
+    EXPECT_TRUE(chunk_generator.next());
+    EXPECT_EQ(64, chunk_generator.addr());
+    EXPECT_EQ(8, chunk_generator.size());
+
+    chunk_generator.setNext(2048);
+    EXPECT_EQ(1024 - 64, chunk_generator.size());
+    EXPECT_TRUE(chunk_generator.last());
+    EXPECT_FALSE(chunk_generator.next());
+    EXPECT_TRUE(chunk_generator.done());
+}
+
 /*
  * A test to consume chunks until the last chunk.
  */
-- 
2.30.2