Fix up removeAll and test with more complex hierarchy

2 files changed, 95 insertions, 6 deletions

diff --git a/lib/collection.h b/lib/collection.h
index 853b982..5f39e71 100644
--- a/lib/collection.h
+++ b/lib/collection.h
@@ -115,14 +115,24 @@ public:
 	}
 
 	template<typename T>
-		requires(std::is_convertible_v<T *, Others *> || ...)
+		requires std::is_base_of_v<Object, T>
 	auto
 	removeAll()
 	{
-		std::get<OtherObjects<T>>(otherObjects).clear();
-		return std::erase_if(objects, [](auto && op) {
-			return dynamic_cast<T *>(op.get());
-		});
+		auto removeAllFrom = [](auto & container) {
+			if constexpr (std::is_base_of_v<T, decltype(std::to_address(container.front()))>) {
+				const auto size = container.size();
+				container.clear();
+				return size;
+			}
+			else {
+				return std::erase_if(container, [](auto && objPtr) -> bool {
+					return dynamic_cast<const T *>(std::to_address(objPtr));
+				});
+			}
+		};
+		(removeAllFrom(std::get<OtherObjects<Others>>(otherObjects)), ...);
+		return removeAllFrom(objects);
 	}
 
 	void
diff --git a/test/test-collection.cpp b/test/test-collection.cpp
index 90a3bd1..b0f2c43 100644
--- a/test/test-collection.cpp
+++ b/test/test-collection.cpp
@@ -42,6 +42,15 @@ public:
 
 class Sub1 : public Sub { };
 
+class Sub2 : public Sub { };
+
+class Base2 {
+public:
+	virtual ~Base2() = default;
+};
+
+class Multi : public Sub1, public Base2 { };
+
 using TestCollection = SharedCollection<Base, Sub>;
 
 BOOST_TEST_DONT_PRINT_LOG_VALUE(TestCollection::Objects::iterator)
@@ -238,8 +247,9 @@ BOOST_AUTO_TEST_CASE(removeAllOfSub)
 	create<Sub>();
 	emplace(std::make_unique<Base>());
 	emplace(std::make_unique<Sub>());
+	emplace(std::make_unique<Sub1>());
 
-	removeAll<Sub>();
+	BOOST_CHECK_EQUAL(removeAll<Sub>(), 3);
 	BOOST_CHECK_EQUAL(objects.size(), 2);
 	BOOST_CHECK(std::get<OtherObjects<Sub>>(otherObjects).empty());
 }
@@ -280,6 +290,75 @@ BOOST_AUTO_TEST_CASE(applyOneType)
 
 BOOST_AUTO_TEST_SUITE_END()
 
+using MultiCollection = Collection<std::unique_ptr<Base>, Multi, Sub, Base2>;
+
+BOOST_FIXTURE_TEST_SUITE(multi, MultiCollection)
+
+BOOST_AUTO_TEST_CASE(addMulti)
+{
+	static_assert(MultiCollection::idx<Multi>() == 0);
+	static_assert(MultiCollection::idx<Sub>() == 1);
+	static_assert(MultiCollection::idx<Base2>() == 2);
+	create<Base>();
+	BOOST_CHECK_EQUAL(objects.size(), 1);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+	create<Sub>();
+	BOOST_CHECK_EQUAL(objects.size(), 2);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 1);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+	create<Sub1>();
+	BOOST_CHECK_EQUAL(objects.size(), 3);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 2);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+	create<Sub2>();
+	BOOST_CHECK_EQUAL(objects.size(), 4);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 3);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+	create<Multi>();
+	BOOST_CHECK_EQUAL(objects.size(), 5);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 1);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 4);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 1);
+}
+
+BOOST_AUTO_TEST_CASE(removeMulti)
+{
+	create<Base>();
+	create<Sub>();
+	create<Sub1>();
+	create<Sub2>();
+	create<Multi>();
+	BOOST_CHECK_EQUAL(objects.size(), 5);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 1);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 4);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 1);
+
+	BOOST_CHECK_EQUAL(removeAll<Multi>(), 1);
+	BOOST_CHECK_EQUAL(objects.size(), 4);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 3);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+
+	BOOST_CHECK_EQUAL(removeAll<Sub>(), 3);
+	BOOST_CHECK_EQUAL(objects.size(), 1);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+
+	BOOST_CHECK_EQUAL(removeAll<Base>(), 1);
+	BOOST_CHECK_EQUAL(objects.size(), 0);
+	BOOST_CHECK_EQUAL(std::get<0>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<1>(otherObjects).size(), 0);
+	BOOST_CHECK_EQUAL(std::get<2>(otherObjects).size(), 0);
+}
+
+BOOST_AUTO_TEST_SUITE_END()
+
 BOOST_AUTO_TEST_CASE(wrapped_ptr_file_cons)
 {
 	using FilePtr = wrapped_ptr<FILE, &fclose>;