From 476822e6c4abf2b1f3d212e6296c1d9d411ff8e8 Mon Sep 17 00:00:00 2001
From: Michele Caini <michele.caini@gmail.com>
Date: Tue, 28 Jan 2025 15:11:16 +0100
Subject: [PATCH] sparse_set: improved shrink_to_fit to also cleanup the sparse
 array

---
 TODO                            |  1 -
 src/entt/entity/sparse_set.hpp  | 26 ++++++++++++++++++++++++++
 test/entt/entity/sparse_set.cpp | 10 +++++++++-
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/TODO b/TODO
index 7971ebcf7..8d28dcd3c 100644
--- a/TODO
+++ b/TODO
@@ -35,7 +35,6 @@ TODO:
 * don't pass reactive storage by default to callback
 * runtime types support for meta for types that aren't backed by C++ types
 * built-in no-pagination storage - no_pagination page size as limits::max
-* sparse_set shrink_to_fit argument for sparse array shrink policy (none, empty, deep, whatever)
 * any cdynamic to support const ownership construction
 * return meta context from meta objects
 * allow passing arguments to meta setter/getter (we can fallback on meta invoke for everything probably)
diff --git a/src/entt/entity/sparse_set.hpp b/src/entt/entity/sparse_set.hpp
index 425af6be3..205aaeb63 100644
--- a/src/entt/entity/sparse_set.hpp
+++ b/src/entt/entity/sparse_set.hpp
@@ -556,6 +556,32 @@ class basic_sparse_set {
 
     /*! @brief Requests the removal of unused capacity. */
     virtual void shrink_to_fit() {
+        sparse_container_type other{sparse.get_allocator()};
+        const auto len = sparse.size();
+        size_type cnt{};
+
+        other.reserve(len);
+
+        for(auto &&elem: std::as_const(packed)) {
+            // also skip tombstones, if any
+            if(const auto page = pos_to_page(entity_to_pos(elem)); page < len && sparse[page] != nullptr) {
+                if(const auto sz = page + 1u; sz > other.size()) {
+                    other.resize(sz, nullptr);
+                }
+
+                other[page] = std::exchange(sparse[page], nullptr);
+
+                if(++cnt == len) {
+                    // early exit due to lack of pages
+                    break;
+                }
+            }
+        }
+
+        release_sparse_pages();
+        sparse.swap(other);
+
+        sparse.shrink_to_fit();
         packed.shrink_to_fit();
     }
 
diff --git a/test/entt/entity/sparse_set.cpp b/test/entt/entity/sparse_set.cpp
index 595376315..f80d6d228 100644
--- a/test/entt/entity/sparse_set.cpp
+++ b/test/entt/entity/sparse_set.cpp
@@ -322,7 +322,15 @@ TYPED_TEST(SparseSet, Pagination) {
 
         set.shrink_to_fit();
 
-        ASSERT_EQ(set.extent(), 2 * traits_type::page_size);
+        switch(policy) {
+        case entt::deletion_policy::swap_and_pop:
+        case entt::deletion_policy::in_place: {
+            ASSERT_EQ(set.extent(), 0u);
+        } break;
+        case entt::deletion_policy::swap_only: {
+            ASSERT_EQ(set.extent(), 2 * traits_type::page_size);
+        } break;
+        }
     }
 }