From af241b031b9e1b112417869bc7149b29d3477fb4 Mon Sep 17 00:00:00 2001
From: Yat Ho <lagoho7@gmail.com>
Date: Sat, 9 Mar 2024 15:27:21 +0800
Subject: [PATCH] fix: destruct elements no longer in use only when erasing

---
 include/small/vector.hpp              | 28 ++++++++++++++---------
 test/unit/CMakeLists.txt              |  1 +
 test/unit/shared_ptr_small_vector.cpp | 33 +++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 11 deletions(-)
 create mode 100644 test/unit/shared_ptr_small_vector.cpp

diff --git a/include/small/vector.hpp b/include/small/vector.hpp
index ba17ba0..feb5252 100644
--- a/include/small/vector.hpp
+++ b/include/small/vector.hpp
@@ -16,6 +16,7 @@
 #include <small/detail/traits/enable_allocator_from_this.hpp>
 #include <small/detail/traits/is_range.hpp>
 #include <small/detail/traits/is_relocatable.hpp>
+#include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdlib>
@@ -1068,33 +1069,38 @@ namespace small {
             if (first == last) {
                 return unconst(first);
             }
-            
-            // Directly destroy elements before mem moving
-            if constexpr (!std::is_trivially_destructible_v<T>) {
-                for (auto it = first; it != last; ++it) {
-                    it->~value_type();
-                }
-            }
-            
+
+            const auto n_erase = last - first;
             if constexpr (is_relocatable_v<value_type> && using_std_allocator) {
                 // Move elements directly in memory
-                const auto n_erase = last - first;
                 const auto n_after_erase = cend() - last;
                 if (n_erase >= n_after_erase) {
                     std::memcpy(
                         (void *) first.base(),
                         (void *) last.base(),
-                        (cend() - last) * sizeof(T));
+                        (cend() - last) * sizeof(value_type));
                 } else {
                     std::memmove(
                         (void *) first.base(),
                         (void *) last.base(),
-                        (cend() - last) * sizeof(T));
+                        (cend() - last) * sizeof(value_type));
                 }
             } else {
                 // Move elements in memory
                 std::move(unconst(last), end(), unconst(first));
+
+                // Destruct elements that were moved from and no longer in-use
+                // N.B. Only do this for non-relocatable types, otherwise you'd
+                // be running the destructor on exact byte copies of in-use
+                // elements, and you might free their internal buffers (oh no!).
+                if constexpr (!std::is_trivially_destructible_v<value_type>) {
+                    std::for_each_n(
+                        crbegin(),
+                        n_erase,
+                        [](value_type const &ele) { ele.~value_type(); });
+                }
             }
+
             // Directly set internal size. Elements are already destroyed.
             set_internal_size(size() - std::distance(first, last));
             return unconst(first);
diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt
index 7de3ad9..6a96fa9 100644
--- a/test/unit/CMakeLists.txt
+++ b/test/unit/CMakeLists.txt
@@ -83,6 +83,7 @@ add_small_test(ptr_wrapper small::small)
 add_small_test(pod_small_vector small::small)
 add_small_test(string_small_vector small::small)
 add_small_test(custom_small_vector small::small)
+add_small_test(shared_ptr_small_vector small::small)
 
 add_small_test(unicode_functions small::small)
 add_small_test(small_string_make small::small)
diff --git a/test/unit/shared_ptr_small_vector.cpp b/test/unit/shared_ptr_small_vector.cpp
new file mode 100644
index 0000000..5f60163
--- /dev/null
+++ b/test/unit/shared_ptr_small_vector.cpp
@@ -0,0 +1,33 @@
+//
+// Copyright (c) 2022 Yat Ho (lagoho7@gmail.com)
+//
+// Distributed under the Boost Software License, Version 1.0.
+// https://www.boost.org/LICENSE_1_0.txt
+//
+
+// C++
+#include <memory>
+
+// External
+#include <catch2/catch.hpp>
+
+// Small
+#include <small/vector.hpp>
+
+TEST_CASE("Shared Ptr Vector") {
+    using namespace small;
+
+    STATIC_REQUIRE(!is_relocatable_v<std::shared_ptr<int>>);
+
+    SECTION("Erase in middle") {
+        vector<std::shared_ptr<int>> a;
+
+        for (int i = 0; i < 2; ++i) {
+            a.emplace_back(std::make_shared<int>(i));
+        }
+
+        a.erase(a.begin());
+
+        REQUIRE(*a[0] == 1);
+    }
+}