Add a ObjectPoolIterator::Type class to allow begin/end iteration on active objects in memory pools.

src/lib/support/Pool.cpp

@@ -127,6 +127,46 @@ Loop StaticAllocatorBitmap::ForEachActiveObjectInner(void * context, Lambda lamb
     return Loop::Finish;
 }
 
+size_t StaticAllocatorBitmap::FirstAllocatedIndex()
+{
+    size_t idx = 0;
+    for (size_t word = 0; word * kBitChunkSize < Capacity(); ++word)
+    {
+        auto & usage = mUsage[word];
+        auto value   = usage.load(std::memory_order_relaxed);
+        for (size_t offset = 0; offset < kBitChunkSize && offset + word * kBitChunkSize < Capacity(); ++offset)
+        {
+            if ((value & (kBit1 << offset)) != 0)
+            {
+                return idx;
+            }
+            idx++;
+        }
+    }
+    VerifyOrDie(idx == mCapacity);
+    return mCapacity;
+}
+
+size_t StaticAllocatorBitmap::NextActiveIndexAfter(size_t start)
+{
+    size_t idx = 0;
+    for (size_t word = 0; word * kBitChunkSize < Capacity(); ++word)
+    {
+        auto & usage = mUsage[word];
+        auto value   = usage.load(std::memory_order_relaxed);
+        for (size_t offset = 0; offset < kBitChunkSize && offset + word * kBitChunkSize < Capacity(); ++offset)
+        {
+            if (((value & (kBit1 << offset)) != 0) && (start < idx))
+            {
+                return idx;
+            }
+            idx++;
+        }
+    }
+    VerifyOrDie(idx == mCapacity);
+    return mCapacity;
+}
+
 #if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
 
 HeapObjectListNode * HeapObjectList::FindNode(void * object) const
@@ -159,24 +199,30 @@ Loop HeapObjectList::ForEachNode(void * context, Lambda lambda)
         p = p->mNext;
     }
     --mIterationDepth;
-    if (mIterationDepth == 0 && mHaveDeferredNodeRemovals)
+    CleanupDeferredReleases();
+    return result;
+}
+
+void HeapObjectList::CleanupDeferredReleases()
+{
+    if (mIterationDepth != 0 || !mHaveDeferredNodeRemovals)
     {
-        // Remove nodes for released objects.
-        p = mNext;
-        while (p != this)
+        return;
+    }
+    // Remove nodes for released objects.
+    HeapObjectListNode * p = mNext;
+    while (p != this)
+    {
+        HeapObjectListNode * next = p->mNext;
+        if (p->mObject == nullptr)
         {
-            HeapObjectListNode * next = p->mNext;
-            if (p->mObject == nullptr)
-            {
-                p->Remove();
-                Platform::Delete(p);
-            }
-            p = next;
+            p->Remove();
+            Platform::Delete(p);
         }
-
-        mHaveDeferredNodeRemovals = false;
+        p = next;
     }
-    return result;
+
+    mHaveDeferredNodeRemovals = false;
 }
 
 #endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP

src/lib/support/Pool.h

@@ -36,6 +36,9 @@
 
 namespace chip {
 
+template <class T>
+class BitmapActiveObjectIterator;
+
 namespace internal {
 
 class Statistics
@@ -91,6 +94,16 @@ class StaticAllocatorBitmap : public internal::StaticAllocatorBase
     void * At(size_t index) { return static_cast<uint8_t *>(mElements) + mElementSize * index; }
     size_t IndexOf(void * element);
 
+    /// Returns the first index that is allocated.
+    ///
+    /// If nothing is allocated, this will return mCapacity
+    size_t FirstAllocatedIndex();
+
+    /// Returns the next active index after `start`.
+    ///
+    /// If nothing else allocated, returns mCapacity
+    size_t NextActiveIndexAfter(size_t start);
+
     using Lambda = Loop (*)(void * context, void * object);
     Loop ForEachActiveObjectInner(void * context, Lambda lambda);
     Loop ForEachActiveObjectInner(void * context, Loop lambda(void * context, const void * object)) const
@@ -102,6 +115,10 @@ class StaticAllocatorBitmap : public internal::StaticAllocatorBase
     void * mElements;
     const size_t mElementSize;
     std::atomic<tBitChunkType> * mUsage;
+
+    /// allow accessing direct At() calls
+    template <class T>
+    friend class ::chip::BitmapActiveObjectIterator;
 };
 
 template <typename T, typename Function>
@@ -132,9 +149,9 @@ struct HeapObjectListNode
         mPrev->mNext = mNext;
     }
 
-    void * mObject;
-    HeapObjectListNode * mNext;
-    HeapObjectListNode * mPrev;
+    void * mObject             = nullptr;
+    HeapObjectListNode * mNext = nullptr;
+    HeapObjectListNode * mPrev = nullptr;
 };
 
 struct HeapObjectList : HeapObjectListNode
@@ -158,6 +175,9 @@ struct HeapObjectList : HeapObjectListNode
         return const_cast<HeapObjectList *>(this)->ForEachNode(context, reinterpret_cast<Lambda>(lambda));
     }
 
+    /// Cleans up any deferred releases IFF iteration depth is 0
+    void CleanupDeferredReleases();
+
     size_t mIterationDepth         = 0;
     bool mHaveDeferredNodeRemovals = false;
 };
@@ -166,6 +186,40 @@ struct HeapObjectList : HeapObjectListNode
 
 } // namespace internal
 
+/// Provides iteration over active objects in a Bitmap pool.
+///
+/// The iterator is valid only if the Pool is not changed (no objects are created
+/// or released)/
+template <class T>
+class BitmapActiveObjectIterator
+{
+public:
+    using value_type = T;
+    using pointer    = T *;
+    using reference  = T &;
+
+    explicit BitmapActiveObjectIterator(internal::StaticAllocatorBitmap * pool, size_t idx) : mPool(pool), mIndex(idx) {}
+    BitmapActiveObjectIterator() {}
+
+    bool operator==(const BitmapActiveObjectIterator & other) const
+    {
+        return (AtEnd() && other.AtEnd()) || ((mPool == other.mPool) && (mIndex == other.mIndex));
+    }
+    bool operator!=(const BitmapActiveObjectIterator & other) const { return !(*this == other); }
+    BitmapActiveObjectIterator & operator++()
+    {
+        mIndex = mPool->NextActiveIndexAfter(mIndex);
+        return *this;
+    }
+    T * operator*() const { return static_cast<T *>(mPool->At(mIndex)); }
+
+private:
+    bool AtEnd() const { return (mPool == nullptr) || (mIndex >= mPool->Capacity()); }
+
+    internal::StaticAllocatorBitmap * mPool = nullptr; // pool that this belongs to
+    size_t mIndex                           = std::numeric_limits<size_t>::max();
+};
+
 /**
  * @class ObjectPool
  *
@@ -205,6 +259,9 @@ class BitMapObjectPool : public internal::StaticAllocatorBitmap
     BitMapObjectPool() : StaticAllocatorBitmap(mData.mMemory, mUsage, N, sizeof(T)) {}
     ~BitMapObjectPool() { VerifyOrDie(Allocated() == 0); }
 
+    BitmapActiveObjectIterator<T> begin() { return BitmapActiveObjectIterator<T>(this, FirstAllocatedIndex()); }
+    BitmapActiveObjectIterator<T> end() { return BitmapActiveObjectIterator<T>(this, N); }
+
     template <typename... Args>
     T * CreateObject(Args &&... args)
     {
@@ -331,6 +388,86 @@ class HeapObjectPool : public internal::Statistics, public HeapObjectPoolExitHan
 #endif // __SANITIZE_ADDRESS__
     }
 
+    /// Provides iteration over active objects in the pool.
+    ///
+    /// NOTE: There is extra logic to allow objects release WHILE the iterator is
+    ///       active while still allowing to advance the iterator.
+    ///       This is done by flagging an iteration depth whenever an active
+    ///       iterator exists. This also means that while a pool iterator exists, releasing
+    ///       of tracking memory objects may be deferred until the last active iterator is
+    ///       released.
+    class ActiveObjectIterator
+    {
+    public:
+        using value_type = T;
+        using pointer    = T *;
+        using reference  = T &;
+
+        ActiveObjectIterator() {}
+        ActiveObjectIterator(const ActiveObjectIterator & other) : mCurrent(other.mCurrent), mEnd(other.mEnd)
+        {
+            if (mEnd != nullptr)
+            {
+                mEnd->mIterationDepth++;
+            }
+        }
+
+        ActiveObjectIterator & operator=(const ActiveObjectIterator & other)
+        {
+            if (mEnd != nullptr)
+            {
+                mEnd->mIterationDepth--;
+                mEnd->CleanupDeferredReleases();
+            }
+            mCurrent = other.mCurrent;
+            mEnd     = other.mEnd;
+            mEnd->mIterationDepth++;
+        }
+
+        ~ActiveObjectIterator()
+        {
+            if (mEnd != nullptr)
+            {
+                mEnd->mIterationDepth--;
+                mEnd->CleanupDeferredReleases();
+            }
+        }
+
+        bool operator==(const ActiveObjectIterator & other) const
+        {
+            // extra current/end compare is to have all "end iterators"
+            // compare as equal (in particular default active object iterator is the end
+            // of an iterator)
+            return (mCurrent == other.mCurrent) || ((mCurrent == mEnd) && (other.mCurrent == other.mEnd));
+        }
+        bool operator!=(const ActiveObjectIterator & other) const { return !(*this == other); }
+        ActiveObjectIterator & operator++()
+        {
+            do
+            {
+                mCurrent = mCurrent->mNext;
+            } while ((mCurrent != mEnd) && (mCurrent->mObject == nullptr));
+            return *this;
+        }
+        T * operator*() const { return static_cast<T *>(mCurrent->mObject); }
+
+    protected:
+        friend class HeapObjectPool<T>;
+
+        explicit ActiveObjectIterator(internal::HeapObjectListNode * current, internal::HeapObjectList * end) :
+            mCurrent(current), mEnd(end)
+        {
+            mEnd->mIterationDepth++;
+        }
+
+    private:
+        internal::HeapObjectListNode * mCurrent = nullptr;
+        internal::HeapObjectList * mEnd         = nullptr;
+    };
+
+    ActiveObjectIterator begin() { return ActiveObjectIterator(mObjects.mNext, &mObjects); }
+    ActiveObjectIterator end() { return ActiveObjectIterator(&mObjects, &mObjects); }
+
     template <typename... Args>
     T * CreateObject(Args &&... args)
     {
@@ -456,6 +593,15 @@ enum class ObjectPoolMem
 #endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
 };
 
+template <typename T, ObjectPoolMem P = ObjectPoolMem::kDefault>
+struct ObjectPoolIterator;
+
+template <typename T>
+struct ObjectPoolIterator<T, ObjectPoolMem::kInline>
+{
+    using Type = BitmapActiveObjectIterator<T>;
+};
+
 template <typename T, size_t N, ObjectPoolMem P = ObjectPoolMem::kDefault>
 class ObjectPool;
 
@@ -465,6 +611,13 @@ class ObjectPool<T, N, ObjectPoolMem::kInline> : public BitMapObjectPool<T, N>
 };
 
 #if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
+template <typename T>
+struct ObjectPoolIterator<T, ObjectPoolMem::kHeap>
+{
+    using Type = typename HeapObjectPool<T>::ActiveObjectIterator;
+};
+
 template <typename T, size_t N>
 class ObjectPool<T, N, ObjectPoolMem::kHeap> : public HeapObjectPool<T>
 {