From 2196db562e7d68bb07ff77b3cb578ebc57956bc1 Mon Sep 17 00:00:00 2001
From: Michele Caini <michele.caini@gmail.com>
Date: Wed, 19 Jan 2022 12:15:22 +0100
Subject: [PATCH] registry:: non-template storage(id) returns an iterator
 rather than a naked pointer

---
 src/entt/entity/registry.hpp  | 39 ++++++++++++++++++-----------------
 test/entt/entity/registry.cpp |  8 +++----
 test/example/entity_copy.cpp  |  4 ++--
 3 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/src/entt/entity/registry.hpp b/src/entt/entity/registry.hpp
index 422e48914..79fcc0b76 100644
--- a/src/entt/entity/registry.hpp
+++ b/src/entt/entity/registry.hpp
@@ -336,6 +336,26 @@ public:
         return iterable_adaptor{internal::storage_proxy_iterator{pools.cbegin()}, internal::storage_proxy_iterator{pools.cend()}};
     }
 
+    /**
+     * @brief Finds the storage associated with a given name, if any.
+     * @param id Name used to map the storage within the registry.
+     * @return An iterator to the given storage if it's found, past the end
+     * iterator otherwise.
+     */
+    [[nodiscard]] auto storage(const id_type id) {
+        return internal::storage_proxy_iterator{pools.find(id)};
+    }
+
+    /**
+     * @brief Finds the storage associated with a given name, if any.
+     * @param id Name used to map the storage within the registry.
+     * @return An iterator to the given storage if it's found, past the end
+     * iterator otherwise.
+     */
+    [[nodiscard]] auto storage(const id_type id) const {
+        return internal::storage_proxy_iterator{pools.find(id)};
+    }
+
     /**
      * @brief Returns the storage for a given component type.
      * @tparam Component Type of component of which to return the storage.
@@ -363,25 +383,6 @@ public:
         return assure<Component>(id);
     }
 
-    /**
-     * @brief Returns the storage associated with a given name, if any.
-     * @param id Name used to map the storage within the registry.
-     * @return The requested storage if it exists, a null pointer otherwise.
-     */
-    [[nodiscard]] base_type *storage(const id_type id) {
-        return const_cast<base_type *>(std::as_const(*this).storage(id));
-    }
-
-    /**
-     * @brief Returns the storage associated with a given name, if any.
-     * @param id Name used to map the storage within the registry.
-     * @return The requested storage if it exists, a null pointer otherwise.
-     */
-    [[nodiscard]] const base_type *storage(const id_type id) const {
-        const auto it = pools.find(id);
-        return it == pools.end() ? nullptr : it->second.get();
-    }
-
     /**
      * @brief Returns the number of entities created so far.
      * @return Number of entities created so far.
diff --git a/test/entt/entity/registry.cpp b/test/entt/entity/registry.cpp
index b38605797..8ab68a5c2 100644
--- a/test/entt/entity/registry.cpp
+++ b/test/entt/entity/registry.cpp
@@ -1881,14 +1881,14 @@ TEST(Registry, RuntimePools) {
     static_assert(std::is_same_v<decltype(registry.storage<empty_type>()), typename entt::storage_traits<entt::entity, empty_type>::storage_type &>);
     static_assert(std::is_same_v<decltype(std::as_const(registry).storage<empty_type>()), const typename entt::storage_traits<entt::entity, empty_type>::storage_type &>);
 
-    static_assert(std::is_same_v<decltype(registry.storage("other"_hs)), typename entt::storage_traits<entt::entity, empty_type>::storage_type::base_type *>);
-    static_assert(std::is_same_v<decltype(std::as_const(registry).storage("other"_hs)), const typename entt::storage_traits<entt::entity, empty_type>::storage_type::base_type *>);
+    static_assert(std::is_same_v<decltype(registry.storage("other"_hs)->second), typename entt::storage_traits<entt::entity, empty_type>::storage_type::base_type &>);
+    static_assert(std::is_same_v<decltype(std::as_const(registry).storage("other"_hs)->second), const typename entt::storage_traits<entt::entity, empty_type>::storage_type::base_type &>);
 
     ASSERT_DEATH([[maybe_unused]] auto &&unused = registry.storage<int>("other"_hs), "");
     ASSERT_DEATH([[maybe_unused]] auto &&unused = std::as_const(registry).storage<int>("other"_hs), "");
 
-    ASSERT_EQ(registry.storage("other"_hs), &storage);
-    ASSERT_EQ(std::as_const(registry).storage("rehto"_hs), nullptr);
+    ASSERT_NE(registry.storage("other"_hs), registry.storage().end());
+    ASSERT_EQ(std::as_const(registry).storage("rehto"_hs), registry.storage().end());
 
     ASSERT_EQ(&registry.storage<empty_type>("other"_hs), &storage);
     ASSERT_NE(&std::as_const(registry).storage<empty_type>(), &storage);
diff --git a/test/example/entity_copy.cpp b/test/example/entity_copy.cpp
index e6cab80f2..254cb89dc 100644
--- a/test/example/entity_copy.cpp
+++ b/test/example/entity_copy.cpp
@@ -56,8 +56,8 @@ TEST(Example, DifferentRegistryTypes) {
     registry.emplace<char>(src, 'c');
 
     for(auto [id, storage]: registry.storage()) {
-        if(auto *pool = other.storage(id); pool && storage.contains(src)) {
-            pool->emplace(dst, storage.get(src));
+        if(auto it = other.storage(id); it != other.storage().end() && storage.contains(src)) {
+            it->second.emplace(dst, storage.get(src));
         }
     }