From 1a34687731859f348618161e275e39068d54164d Mon Sep 17 00:00:00 2001
From: "Bernhart, Bryan" <bryan.bernhart@intel.com>
Date: Mon, 11 Sep 2023 11:03:30 -0700
Subject: [PATCH] Refactor pools to iterate as STL containers.

---
 src/gpgmm/common/IndexedMemoryPool.cpp | 15 ++++++++++++++-
 src/gpgmm/common/IndexedMemoryPool.h   | 11 +++++++++--
 src/gpgmm/common/LIFOMemoryPool.cpp    | 15 ++++++++++++++-
 src/gpgmm/common/LIFOMemoryPool.h      | 11 +++++++++--
 src/gpgmm/common/MemoryPool.h          |  9 +++++----
 5 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/src/gpgmm/common/IndexedMemoryPool.cpp b/src/gpgmm/common/IndexedMemoryPool.cpp
index 6ef66218..fcf360b9 100644
--- a/src/gpgmm/common/IndexedMemoryPool.cpp
+++ b/src/gpgmm/common/IndexedMemoryPool.cpp
@@ -36,7 +36,7 @@ namespace gpgmm {
     }
 
     uint64_t IndexedMemoryPool::ReleasePool(uint64_t bytesToRelease) {
-        return TrimPoolUntil(mPool, bytesToRelease);
+        return TrimPoolUntil(this, bytesToRelease);
     }
 
     uint64_t IndexedMemoryPool::GetPoolSize() const {
@@ -48,4 +48,17 @@ namespace gpgmm {
         }
         return count;
     }
+
+    IndexedMemoryPool::UnderlyingContainerType::iterator IndexedMemoryPool::begin() {
+        return mPool.begin();
+    }
+
+    IndexedMemoryPool::UnderlyingContainerType::iterator IndexedMemoryPool::end() {
+        return mPool.end();
+    }
+
+    void IndexedMemoryPool::ShrinkPool(uint64_t lastIndex) {
+        mPool.erase(begin(), begin() + lastIndex);
+    }
+
 }  // namespace gpgmm
diff --git a/src/gpgmm/common/IndexedMemoryPool.h b/src/gpgmm/common/IndexedMemoryPool.h
index 4badfb58..dc62d518 100644
--- a/src/gpgmm/common/IndexedMemoryPool.h
+++ b/src/gpgmm/common/IndexedMemoryPool.h
@@ -22,6 +22,8 @@
 namespace gpgmm {
 
     class IndexedMemoryPool final : public MemoryPoolBase {
+        using UnderlyingContainerType = std::vector<std::unique_ptr<MemoryAllocationBase>>;
+
       public:
         explicit IndexedMemoryPool(uint64_t memorySize);
         ~IndexedMemoryPool() override = default;
@@ -31,11 +33,16 @@ namespace gpgmm {
         void ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
                           uint64_t indexInPool) override;
         uint64_t ReleasePool(uint64_t bytesToRelease) override;
-
         uint64_t GetPoolSize() const override;
 
+        UnderlyingContainerType::iterator begin();
+        UnderlyingContainerType::iterator end();
+
+        // Resizes the pool up to but not including |lastIndex|.
+        void ShrinkPool(uint64_t lastIndex);
+
       private:
-        std::vector<std::unique_ptr<MemoryAllocationBase>> mPool;
+        UnderlyingContainerType mPool;
     };
 
 }  // namespace gpgmm
diff --git a/src/gpgmm/common/LIFOMemoryPool.cpp b/src/gpgmm/common/LIFOMemoryPool.cpp
index 7032cf56..cfa831e2 100644
--- a/src/gpgmm/common/LIFOMemoryPool.cpp
+++ b/src/gpgmm/common/LIFOMemoryPool.cpp
@@ -42,10 +42,23 @@ namespace gpgmm {
     }
 
     uint64_t LIFOMemoryPool::ReleasePool(uint64_t bytesToRelease) {
-        return TrimPoolUntil(mPool, bytesToRelease);
+        return TrimPoolUntil(this, bytesToRelease);
     }
 
     uint64_t LIFOMemoryPool::GetPoolSize() const {
         return mPool.size();
     }
+
+    LIFOMemoryPool::UnderlyingContainerType::iterator LIFOMemoryPool::begin() {
+        return mPool.begin();
+    }
+
+    LIFOMemoryPool::UnderlyingContainerType::iterator LIFOMemoryPool::end() {
+        return mPool.end();
+    }
+
+    void LIFOMemoryPool::ShrinkPool(uint64_t lastIndex) {
+        mPool.erase(begin(), begin() + lastIndex);
+    }
+
 }  // namespace gpgmm
diff --git a/src/gpgmm/common/LIFOMemoryPool.h b/src/gpgmm/common/LIFOMemoryPool.h
index 68227162..2993cba2 100644
--- a/src/gpgmm/common/LIFOMemoryPool.h
+++ b/src/gpgmm/common/LIFOMemoryPool.h
@@ -23,6 +23,8 @@ namespace gpgmm {
 
     // Pool using LIFO (newest are recycled first).
     class LIFOMemoryPool : public MemoryPoolBase {
+        using UnderlyingContainerType = std::deque<std::unique_ptr<MemoryAllocationBase>>;
+
       public:
         explicit LIFOMemoryPool(uint64_t memorySize);
         ~LIFOMemoryPool() override = default;
@@ -33,11 +35,16 @@ namespace gpgmm {
         void ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
                           uint64_t indexInPool = kInvalidIndex) override;
         uint64_t ReleasePool(uint64_t bytesToFree = kInvalidSize) override;
-
         uint64_t GetPoolSize() const override;
 
+        UnderlyingContainerType::iterator begin();
+        UnderlyingContainerType::iterator end();
+
+        // Resizes the pool up to but not including |lastIndex|.
+        void ShrinkPool(uint64_t lastIndex);
+
       private:
-        std::deque<std::unique_ptr<MemoryAllocationBase>> mPool;
+        UnderlyingContainerType mPool;
     };
 
 }  // namespace gpgmm
diff --git a/src/gpgmm/common/MemoryPool.h b/src/gpgmm/common/MemoryPool.h
index a959d53c..1c8d2090 100644
--- a/src/gpgmm/common/MemoryPool.h
+++ b/src/gpgmm/common/MemoryPool.h
@@ -16,6 +16,7 @@
 #define SRC_GPGMM_COMMON_MEMORYPOOL_H_
 
 #include "gpgmm/common/MemoryAllocation.h"
+#include "gpgmm/common/MemoryAllocator.h"
 #include "gpgmm/utils/Limits.h"
 
 #include <memory>
@@ -49,11 +50,11 @@ namespace gpgmm {
 
       protected:
         // Shrinks the size of the pool in |mMemorySize| sizes until |bytesToRelease| is reached.
-        template <typename T>
-        uint64_t TrimPoolUntil(T& pool, uint64_t bytesToRelease) {
+        template <typename MemoryPoolT>
+        uint64_t TrimPoolUntil(MemoryPoolT* pool, uint64_t bytesToRelease) {
             uint64_t totalBytesReleased = 0;
             uint64_t lastIndex = 0;
-            for (auto& allocation : pool) {
+            for (auto& allocation : *pool) {
                 totalBytesReleased += allocation->GetSize();
                 allocation->GetAllocator()->DeallocateMemory(std::move(allocation));
                 lastIndex++;
@@ -64,7 +65,7 @@ namespace gpgmm {
 
             // Last is non-inclusive or [first, last).
             if (lastIndex > 0) {
-                pool.erase(pool.begin(), pool.begin() + lastIndex);
+                pool->ShrinkPool(lastIndex);
             }
 
             return totalBytesReleased;