BLI: wrap more features off tbb::enumerable_thread_specific

* Make the wrapper enumerable.
* Support an initializer function.

1 files changed, 49 insertions, 2 deletions

diff --git a/source/blender/blenlib/BLI_enumerable_thread_specific.hh b/source/blender/blenlib/BLI_enumerable_thread_specific.hh
index a05f7724dd2..25fd02b41fb 100644
--- a/source/blender/blenlib/BLI_enumerable_thread_specific.hh
+++ b/source/blender/blenlib/BLI_enumerable_thread_specific.hh
@@ -46,25 +46,72 @@ template<typename T> class EnumerableThreadSpecific : NonCopyable, NonMovable {
   tbb::enumerable_thread_specific<T> values_;
 
  public:
+  using iterator = typename tbb::enumerable_thread_specific<T>::iterator;
+
+  EnumerableThreadSpecific() = default;
+
+  template<typename F> EnumerableThreadSpecific(F initializer) : values_(std::move(initializer))
+  {
+  }
+
   T &local()
   {
     return values_.local();
   }
 
+  iterator begin()
+  {
+    return values_.begin();
+  }
+
+  iterator end()
+  {
+    return values_.end();
+  }
+
 #else /* WITH_TBB */
 
  private:
   std::mutex mutex_;
   /* Maps thread ids to their corresponding values. The values are not embedded in the map, so that
    * their addresses do not change when the map grows. */
-  Map<int, std::unique_ptr<T>> values_;
+  Map<int, std::reference_wrapper<T>> values_;
+  Vector<std::unique_ptr<T>> owned_values_;
+  std::function<T()> initializer_;
 
  public:
+  using iterator = typename Map<int, std::reference_wrapper<T>>::MutableValueIterator;
+
+  EnumerableThreadSpecific() : initializer_([]() { return T(); })
+  {
+  }
+
+  template<typename F>
+  EnumerableThreadSpecific(F initializer) : initializer_(std::move(initializer))
+  {
+  }
+
   T &local()
   {
     const int thread_id = enumerable_thread_specific_utils::thread_id;
     std::lock_guard lock{mutex_};
-    return *values_.lookup_or_add_cb(thread_id, []() { return std::make_unique<T>(); });
+    return values_.lookup_or_add_cb(thread_id, [&]() {
+      /* `std::make_unique` does not work here if T is non-copyable and non-movable. */
+      std::unique_ptr<T> value{new T(initializer_())};
+      std::reference_wrapper<T> ref = *value;
+      owned_values_.append(std::move(value));
+      return ref;
+    });
+  }
+
+  iterator begin()
+  {
+    return values_.values().begin();
+  }
+
+  iterator end()
+  {
+    return values_.values().end();
   }
 
 #endif /* WITH_TBB */