[Bf-blender-cvs] [476610dfae1] soc-2021-adaptive-cloth: bli: generational_arena: insert_with() and try_insert_with()

[ishbosamiya]

Commit: 476610dfae14b012fe90795e3ec5b90585a3c5c1
Author: ishbosamiya
Date:   Sat Jun 19 23:54:06 2021 +0530
Branches: soc-2021-adaptive-cloth
https://developer.blender.org/rB476610dfae14b012fe90795e3ec5b90585a3c5c1

bli: generational_arena: insert_with() and try_insert_with()

===================================================================

M	source/blender/blenlib/BLI_generational_arena.hh
M	source/blender/blenlib/tests/BLI_generational_arena_test.cc

===================================================================

diff --git a/source/blender/blenlib/BLI_generational_arena.hh b/source/blender/blenlib/BLI_generational_arena.hh
index 3e9a12450a2..2d588907a06 100644
--- a/source/blender/blenlib/BLI_generational_arena.hh
+++ b/source/blender/blenlib/BLI_generational_arena.hh
@@ -252,6 +252,50 @@ class Arena {
     return Index::invalid();
   }
 
+  /* TODO(ish): add optimization by moving `f`, can be done by
+   * returning value if `try_insert()` fails */
+  std::optional<Index> try_insert_with(std::function<T(Index)> f)
+  {
+    if (this->next_free_head) {
+      auto loc = *this->next_free_head;
+
+      if (auto entry = std::get_if<EntryNoExist>(&this->data[loc])) {
+        this->next_free_head = entry->next_free;
+        Index index(loc, this->generation);
+        T value = f(index);
+        this->data[loc] = EntryExist(value, this->generation);
+        this->length += 1;
+
+        return index;
+      }
+
+      /* The linked list created to
+       * know where to insert next is
+       * corrupted.
+       * `this->next_free_head` is corrupted */
+      BLI_assert_unreachable();
+    }
+    return std::nullopt;
+  }
+
+  Index insert_with(std::function<T(Index)> f)
+  {
+    if (auto index = this->try_insert_with(f)) {
+      return *index;
+    }
+
+    /* couldn't insert the value within reserved memory space */
+    const auto reserve_cap = this->data.size() == 0 ? 1 : this->data.size();
+    this->reserve(reserve_cap * 2);
+    if (auto index = this->try_insert_with(f)) {
+      return *index;
+    }
+
+    /* now that more memory has been reserved, it shouldn't fail */
+    BLI_assert_unreachable();
+    return Index::invalid();
+  }
+
   std::optional<T> remove(Index index)
   {
     if (index.index >= this->data.size()) {
diff --git a/source/blender/blenlib/tests/BLI_generational_arena_test.cc b/source/blender/blenlib/tests/BLI_generational_arena_test.cc
index e677b293434..ead72f80f35 100644
--- a/source/blender/blenlib/tests/BLI_generational_arena_test.cc
+++ b/source/blender/blenlib/tests/BLI_generational_arena_test.cc
@@ -2,6 +2,9 @@
 
 #include "testing/testing.h"
 
+#include <functional>
+#include <tuple>
+
 namespace blender::tests {
 
 using namespace blender::generational_arena;
@@ -56,6 +59,24 @@ TEST(generational_arena, Insert2)
   EXPECT_EQ(arena.get(i5), 5);
 }
 
+TEST(generational_arena, InsertWith)
+{
+  Arena<std::tuple<Index, int>> arena;
+  auto i0 = arena.insert_with([](Index index) { return std::make_tuple(index, 0); });
+  auto i1 = arena.insert_with([](Index index) { return std::make_tuple(index, 1); });
+  auto i2 = arena.insert_with([](Index index) { return std::make_tuple(index, 2); });
+  auto i3 = arena.insert_with([](Index index) { return std::make_tuple(index, 3); });
+  auto i4 = arena.insert_with([](Index index) { return std::make_tuple(index, 4); });
+
+  EXPECT_EQ(arena.capacity(), 8);
+  EXPECT_EQ(arena.size(), 5);
+  EXPECT_EQ(arena.get(i0).value().get(), std::make_tuple(i0, 0));
+  EXPECT_EQ(arena.get(i1).value().get(), std::make_tuple(i1, 1));
+  EXPECT_EQ(arena.get(i2).value().get(), std::make_tuple(i2, 2));
+  EXPECT_EQ(arena.get(i3).value().get(), std::make_tuple(i3, 3));
+  EXPECT_EQ(arena.get(i4).value().get(), std::make_tuple(i4, 4));
+}
+
 TEST(generational_arena, Remove)
 {
   Arena<int> arena(3);