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..3051d980d45 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<void(void *)> initializer_;
 
  public:
+  using iterator = typename Map<int, std::reference_wrapper<T>>::MutableValueIterator;
+
+  EnumerableThreadSpecific() : initializer_([](void *buffer) { new (buffer) T(); })
+  {
+  }
+
+  template<typename F>
+  EnumerableThreadSpecific(F initializer)
+      : initializer_([=](void *buffer) { new (buffer) T(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, [&]() {
+      T *value = (T *)::operator new(sizeof(T));
+      initializer_(value);
+      owned_values_.append(std::unique_ptr<T>{value});
+      return std::reference_wrapper<T>{*value};
+    });
+  }
+
+  iterator begin()
+  {
+    return values_.values().begin();
+  }
+
+  iterator end()
+  {
+    return values_.values().end();
   }
 
 #endif /* WITH_TBB */