registry: more stable assure function (with tests) - close #1148

src/entt/entity/registry.hpp

@@ -249,10 +249,10 @@ class basic_registry {
             return entities;
         } else {
             using storage_type = storage_for_type<Type>;
-            auto &cpool = pools[id];
 
-            if(!cpool) {
+            if(auto it = pools.find(id); it == pools.cend()) {
                 using alloc_type = typename storage_type::allocator_type;
+                typename pool_container_type::mapped_type cpool{};
 
                 if constexpr(std::is_void_v<Type> && !std::is_constructible_v<alloc_type, allocator_type>) {
                     // std::allocator<void> has no cross constructors (waiting for C++20)
@@ -261,11 +261,14 @@ class basic_registry {
                     cpool = std::allocate_shared<storage_type>(get_allocator(), get_allocator());
                 }
 
+                pools.emplace(id, cpool);
                 cpool->bind(forward_as_any(*this));
-            }
 
-            ENTT_ASSERT(cpool->type() == type_id<Type>(), "Unexpected type");
-            return static_cast<storage_type &>(*cpool);
+                return static_cast<storage_type &>(*cpool);
+            } else {
+                ENTT_ASSERT(it->second->type() == type_id<Type>(), "Unexpected type");
+                return static_cast<storage_type &>(*it->second);
+            }
         }
     }
 

test/common/mixin.hpp

@@ -0,0 +1,36 @@
+#ifndef ENTT_COMMON_MIXIN_HPP
+#define ENTT_COMMON_MIXIN_HPP
+
+#include <entt/core/any.hpp>
+#include <entt/entity/fwd.hpp>
+
+namespace test {
+
+struct assure_loop {};
+
+template<typename Type>
+class assure_loop_mixin: public Type {
+    using underlying_type = Type;
+    using registry_type = entt::basic_registry<typename underlying_type::entity_type, typename underlying_type::base_type::allocator_type>;
+
+public:
+    using allocator_type = typename underlying_type::allocator_type;
+    using entity_type = typename underlying_type::entity_type;
+
+    using Type::Type;
+
+    void bind(entt::any value) noexcept override {
+        if(auto *owner = entt::any_cast<registry_type>(&value)) {
+            owner->storage<int>();
+        }
+    }
+};
+
+} // namespace test
+
+template<typename Entity>
+struct entt::storage_type<test::assure_loop, Entity> {
+    using type = test::assure_loop_mixin<entt::basic_storage<test::assure_loop, Entity>>;
+};
+
+#endif

test/entt/entity/registry.cpp

@@ -19,6 +19,7 @@
 #include "../../common/config.h"
 #include "../../common/empty.h"
 #include "../../common/entity.h"
+#include "../../common/mixin.hpp"
 #include "../../common/non_default_constructible.h"
 #include "../../common/pointer_stable.h"
 
@@ -2424,6 +2425,22 @@ TEST(Registry, ScramblingPoolsIsAllowed) {
     });
 }
 
+TEST(Registry, AssureMixinLoop) {
+    entt::registry registry{};
+    const auto entity = registry.create();
+
+    ASSERT_EQ(std::as_const(registry).storage<test::assure_loop>(), nullptr);
+    ASSERT_EQ(std::as_const(registry).storage<int>(), nullptr);
+
+    registry.emplace<test::assure_loop>(entity);
+
+    ASSERT_NE(std::as_const(registry).storage<test::assure_loop>(), nullptr);
+    ASSERT_NE(std::as_const(registry).storage<int>(), nullptr);
+
+    ASSERT_TRUE(registry.all_of<test::assure_loop>(entity));
+    ASSERT_FALSE(registry.all_of<int>(entity));
+}
+
 TEST(Registry, VoidType) {
     using namespace entt::literals;