1 files changed, 63 insertions, 32 deletions

diff --git a/src/cppc4group.cpp b/src/cppc4group.cpp
index 329d6de..b3bbd42 100644
--- a/src/cppc4group.cpp
+++ b/src/cppc4group.cpp
@@ -40,19 +40,74 @@ struct CppC4Group::Private {
 };
 
 struct CppC4Group::MemoryManagement::Private {
-	bool predefined;
+	struct CustomMemoryManagement {
+		CC4Group_MemoryManagement_t realMemoryManagement;
+
+		Start start;
+		End end;
+
+		size_t referenceCount;
+
+		static void* callStart(void* const data, size_t const size, void* const arg)
+		{
+			return reinterpret_cast<CustomMemoryManagement*>(arg)->start(data, size);
+		}
+
+		static void callEndAndCleanup(void* const data, void* const arg)
+		{
+			auto customMemoryManagement = reinterpret_cast<CustomMemoryManagement*>(arg);
+			customMemoryManagement->end(data);
+			customMemoryManagement->dereference();
+		}
+
+		CustomMemoryManagement(Start&& start, End&& end) : realMemoryManagement{callStart, callEndAndCleanup, static_cast<void*>(this)}, start{start}, end{end}, referenceCount{}
+		{
+
+		}
+
+		void reference()
+		{
+			++referenceCount;
+		}
+
+		void dereference()
+		{
+			if(--referenceCount == 0)
+			{
+				delete this;
+			}
+		}
+	};
+
 	MemoryManagementStrategy strategy;
+	CustomMemoryManagement* customStrategyData;
+
+	Private(MemoryManagementStrategy strategy) : strategy{strategy}, customStrategyData{}
+	{
+
+	}
+
+	Private(Start&& start, End&& end) : customStrategyData{new CustomMemoryManagement{std::forward<Start>(start), std::forward<End>(end)}}
+	{
+		customStrategyData->reference();
+		strategy = reinterpret_cast<MemoryManagementStrategy>(&customStrategyData->realMemoryManagement);
+	}
 
-	Start start;
-	End end;
+	~Private()
+	{
+		if(customStrategyData != nullptr)
+		{
+			customStrategyData->dereference();
+		}
+	}
 };
 
-CppC4Group::MemoryManagement::MemoryManagement(MemoryManagementStrategy strategy) : p{new Private{true, strategy}}
+CppC4Group::MemoryManagement::MemoryManagement(MemoryManagementStrategy strategy) : p{new Private{strategy}}
 {
 
 }
 
-CppC4Group::MemoryManagement::MemoryManagement(Start&& start, End&& end) : p{new Private{false, {}, start, end}}
+CppC4Group::MemoryManagement::MemoryManagement(Start&& start, End&& end) : p{new Private{std::forward<Start>(start), std::forward<End>(end)}}
 {
 
 }
@@ -64,35 +119,11 @@ CppC4Group::MemoryManagement::~MemoryManagement()
 
 CppC4Group::MemoryManagement::MemoryManagementStrategy CppC4Group::MemoryManagement::operator()() const
 {
-	if(p->predefined)
+	if(p->customStrategyData != nullptr)
 	{
-		return p->strategy;
-	}
-	else
-	{
-		struct CustomMemoryManagement {
-			CC4Group_MemoryManagement_t realMemoryManagement;
-
-			Start start;
-			End end;
-
-			static void* callStart(void* const data, size_t const size, void* const arg)
-			{
-				return reinterpret_cast<CustomMemoryManagement*>(arg)->start(data, size);
-			}
-
-			static void callEndAndCleanup(void* const data, void* const arg)
-			{
-				auto customMemoryManagement = reinterpret_cast<CustomMemoryManagement*>(arg);
-				customMemoryManagement->end(data);
-				delete customMemoryManagement;
-			}
-		};
-
-		auto customManagement = new CustomMemoryManagement{{CustomMemoryManagement::callStart, CustomMemoryManagement::callEndAndCleanup}, p->start, p->end};
-		customManagement->realMemoryManagement.arg = static_cast<void*>(customManagement);
-		return reinterpret_cast<MemoryManagementStrategy>(&customManagement->realMemoryManagement);
+		p->customStrategyData->reference();
 	}
+	return p->strategy;
 }
 
 const CppC4Group::MemoryManagement CppC4Group::MemoryManagement::Take{reinterpret_cast<MemoryManagement::MemoryManagementStrategy>(cc4group.MemoryManagement.Take)};