bli: generational_arena: add iterator support

Made as a bidirectional iterator since movement can be in both
directions. Random access iterator is not possible since there can be
gaps in between the elements.

Had a very annoying bug- using the iterator with the standard library
would throw an error "`iterator_category` not defined" even though it
was defined. Spent a lot of time to realize that `difference_type`
should also be defined but other types within the iterator class are
optional. For some reason the compiler does not create the
`iterator_trait` correctly if `difference_type` is missing but throws
error about `iterator_category` which is extremely strange and time consuming to debug :(

2 files changed, 197 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_generational_arena.hh b/source/blender/blenlib/BLI_generational_arena.hh
index 2d588907a06..84c2b391445 100644
--- a/source/blender/blenlib/BLI_generational_arena.hh
+++ b/source/blender/blenlib/BLI_generational_arena.hh
@@ -52,12 +52,15 @@
  */
 /* TODO(ish): need to complete documentation */
 
+#include <cstddef>
 #include <functional>
+#include <iterator>
 #include <limits>
 #include <optional>
 #include <tuple>
 #include <variant>
 
+#include "BLI_assert.h"
 #include "BLI_vector.hh"
 
 #include "testing/testing.h"
@@ -125,6 +128,10 @@ template<
      */
     typename Allocator = GuardedAllocator>
 class Arena {
+ public:
+  class Iterator;
+
+ private:
   struct EntryNoExist;
   struct EntryExist;
   /* using declarations */
@@ -406,6 +413,112 @@ class Arena {
     return static_cast<isize>(this->length);
   }
 
+  Iterator begin()
+  {
+    return Iterator(this->data.begin(), this->data.begin(), this->data.end());
+  }
+
+  Iterator end()
+  {
+    return Iterator(this->data.end(), this->data.begin(), this->data.end());
+  }
+
+  class Iterator {
+   public:
+    using iterator_category = std::bidirectional_iterator_tag;
+    using difference_type = std::ptrdiff_t;
+    using value_type = T;
+    using pointer = value_type *;
+    using reference = value_type &;
+
+   private:
+    Entry *ptr;   /* points to current position */
+    Entry *start; /* points to first element in the
+                   * Arena::data, aka Arena::data.begin() */
+    Entry *end;   /* points to last+1 element in the Arena::data, aka Arena::data.end()*/
+
+   public:
+    Iterator(Entry *ptr, Entry *start, Entry *end) : ptr(ptr), start(start), end(end)
+    {
+    }
+
+    reference operator*() const
+    {
+      if (auto val = std::get_if<EntryExist>(this->ptr)) {
+        return val->value;
+      }
+
+      BLI_assert_unreachable();
+
+      return std::get<EntryExist>(*this->ptr).value;
+    }
+
+    pointer operator->()
+    {
+      return this->ptr;
+    }
+
+    /* pre fix */
+    Iterator &operator++()
+    {
+      BLI_assert(this->ptr != this->end);
+      while (true) {
+        this->ptr++;
+
+        if (this->ptr == this->end) {
+          break;
+        }
+
+        if (auto val = std::get_if<EntryExist>(this->ptr)) {
+          break;
+        }
+      }
+      return *this;
+    }
+
+    Iterator &operator--()
+    {
+      BLI_assert(this->ptr != this->start);
+      while (true) {
+        this->ptr--;
+
+        if (this->ptr == this->start) {
+          break;
+        }
+
+        if (auto val = std::get_if<EntryExist>(this->ptr)) {
+          break;
+        }
+      }
+      return *this;
+    }
+
+    /* post fix */
+    Iterator operator++(int)
+    {
+      Iterator temp = *this;
+      ++(*this);
+      return temp;
+    }
+
+    Iterator operator--(int)
+    {
+      Iterator temp = *this;
+      --(*this);
+      return temp;
+    }
+
+    friend bool operator==(const Iterator &a, const Iterator &b)
+    {
+      return a.start == b.start && a.end == b.end && a.ptr == b.ptr;
+    }
+
+    friend bool operator!=(const Iterator &a, const Iterator &b)
+    {
+      return a.start != b.start || a.end != b.end || a.ptr != b.ptr;
+    }
+  };
+
  protected:
   /* all protected static methods */
   /* all protected non-static methods */
diff --git a/source/blender/blenlib/tests/BLI_generational_arena_test.cc b/source/blender/blenlib/tests/BLI_generational_arena_test.cc
index ead72f80f35..279480dbb44 100644
--- a/source/blender/blenlib/tests/BLI_generational_arena_test.cc
+++ b/source/blender/blenlib/tests/BLI_generational_arena_test.cc
@@ -2,7 +2,9 @@
 
 #include "testing/testing.h"
 
+#include <algorithm>
 #include <functional>
+#include <gtest/gtest.h>
 #include <tuple>
 
 namespace blender::tests {
@@ -145,6 +147,88 @@ TEST(generational_arena, GetNoGenIndex)
   EXPECT_EQ(arena.get_no_gen(2), std::nullopt);
 }
 
+TEST(generational_arena, Iter)
+{
+  Arena<int> arena;
+  arena.insert(0);
+  arena.insert(0);
+  arena.insert(0);
+  arena.insert(0);
+  arena.insert(0);
+
+  for (const auto &i : arena) {
+    EXPECT_EQ(i, 0);
+  }
+}
+
+TEST(generational_arena, Iter2)
+{
+  Arena<int> arena;
+  arena.insert(2);
+  arena.insert(1);
+  arena.insert(4);
+  arena.insert(3);
+  arena.insert(0);
+
+  EXPECT_TRUE(std::any_of(arena.begin(), arena.end(), [](const int &val) { return val % 2; }));
+
+  auto it = std::partition(arena.begin(), arena.end(), [](const int &val) { return val % 2; });
+
+  EXPECT_NE(std::find(arena.begin(), it, 1), arena.end());
+  EXPECT_NE(std::find(arena.begin(), it, 3), arena.end());
+  EXPECT_NE(std::find(it, arena.end(), 0), arena.end());
+  EXPECT_NE(std::find(it, arena.end(), 2), arena.end());
+  EXPECT_NE(std::find(it, arena.end(), 4), arena.end());
+}
+
+TEST(generational_arena, IterIncrement)
+{
+  Arena<int> arena;
+  arena.insert(0);
+  arena.insert(1);
+  auto i2 = arena.insert(2);
+  arena.insert(3);
+  arena.insert(4);
+
+  arena.remove(i2);
+
+  auto iter = arena.begin();
+  EXPECT_EQ(*iter, 0);
+  iter++;
+  EXPECT_EQ(*iter, 1);
+  iter++;
+  EXPECT_EQ(*iter, 3);
+  ++iter;
+  EXPECT_EQ(*iter, 4);
+  ++iter;
+  EXPECT_EQ(iter, arena.end());
+}
+
+TEST(generational_arena, IterDecrement)
+{
+  Arena<int> arena;
+  arena.insert(0);
+  arena.insert(1);
+  auto i2 = arena.insert(2);
+  arena.insert(3);
+  arena.insert(4);
+
+  arena.remove(i2);
+
+  auto iter = arena.end();
+  --iter;
+  EXPECT_EQ(*iter, 4);
+  iter--;
+  EXPECT_EQ(*iter, 3);
+  iter--;
+  EXPECT_EQ(*iter, 1);
+  --iter;
+  EXPECT_EQ(*iter, 0);
+  EXPECT_EQ(iter, arena.begin());
+
+  EXPECT_BLI_ASSERT(--iter, "");
+}
+
 } /* namespace blender::tests */
 
 namespace blender::generational_arena {