Fix catastrophic logic mistake in cppc4group's custom memory management implementation

1 files changed, 20 insertions, 8 deletions

diff --git a/src/cppc4group.cpp b/src/cppc4group.cpp
index b3bbd42..04f7cef 100644
--- a/src/cppc4group.cpp
+++ b/src/cppc4group.cpp
@@ -47,20 +47,23 @@ struct CppC4Group::MemoryManagement::Private {
 		End end;
 
 		size_t referenceCount;
+		size_t usageReferenceCount;
 
 		static void* callStart(void* const data, size_t const size, void* const arg)
 		{
-			return reinterpret_cast<CustomMemoryManagement*>(arg)->start(data, size);
+			auto customMemoryManagement = reinterpret_cast<CustomMemoryManagement*>(arg);
+			customMemoryManagement->referenceUsage();
+			return customMemoryManagement->start(data, size);
 		}
 
 		static void callEndAndCleanup(void* const data, void* const arg)
 		{
 			auto customMemoryManagement = reinterpret_cast<CustomMemoryManagement*>(arg);
 			customMemoryManagement->end(data);
-			customMemoryManagement->dereference();
+			customMemoryManagement->dereferenceUsage();
 		}
 
-		CustomMemoryManagement(Start&& start, End&& end) : realMemoryManagement{callStart, callEndAndCleanup, static_cast<void*>(this)}, start{start}, end{end}, referenceCount{}
+		CustomMemoryManagement(Start&& start, End&& end) : realMemoryManagement{callStart, callEndAndCleanup, static_cast<void*>(this)}, start{start}, end{end}, referenceCount{}, usageReferenceCount{}
 		{
 
 		}
@@ -72,7 +75,20 @@ struct CppC4Group::MemoryManagement::Private {
 
 		void dereference()
 		{
-			if(--referenceCount == 0)
+			if(--referenceCount == 0 && usageReferenceCount == 0)
+			{
+				delete this;
+			}
+		}
+
+		void referenceUsage()
+		{
+			++usageReferenceCount;
+		}
+
+		void dereferenceUsage()
+		{
+			if(--usageReferenceCount == 0 && referenceCount == 0)
 			{
 				delete this;
 			}
@@ -119,10 +135,6 @@ CppC4Group::MemoryManagement::~MemoryManagement()
 
 CppC4Group::MemoryManagement::MemoryManagementStrategy CppC4Group::MemoryManagement::operator()() const
 {
-	if(p->customStrategyData != nullptr)
-	{
-		p->customStrategyData->reference();
-	}
 	return p->strategy;
 }