From 8b486ed62573867b807a343cb263e9134f45b5ba Mon Sep 17 00:00:00 2001
From: George FunBook <gkurelic@gmail.com>
Date: Wed, 10 Apr 2024 11:28:19 -0500
Subject: [PATCH] add FLX_TRACK_POOLS for finding pool leaks

---
 flixel/system/macros/FlxDefines.hx |   7 ++
 flixel/util/FlxPool.hx             | 110 +++++++++++++++++++++++++++--
 tests/coverage/Project.xml         |   1 +
 3 files changed, 114 insertions(+), 4 deletions(-)

diff --git a/flixel/system/macros/FlxDefines.hx b/flixel/system/macros/FlxDefines.hx
index 1d167c62d2..5b6e2642e3 100644
--- a/flixel/system/macros/FlxDefines.hx
+++ b/flixel/system/macros/FlxDefines.hx
@@ -39,6 +39,8 @@ private enum UserDefines
 	FLX_NO_POINT_POOL;
 	FLX_NO_PITCH;
 	FLX_NO_SAVE;
+	/** Adds trackers to FlxPool instances, only available on debug*/
+	FLX_TRACK_POOLS;
 }
 
 /**
@@ -80,6 +82,7 @@ private enum HelperDefines
 	FLX_NO_CI;
 	FLX_SAVE;
 	FLX_HEALTH;
+	FLX_NO_TRACK_POOLS;
 }
 
 class FlxDefines
@@ -179,6 +182,7 @@ class FlxDefines
 		defineInversion(FLX_COVERAGE_TEST, FLX_NO_COVERAGE_TEST);
 		defineInversion(FLX_SWF_VERSION_TEST, FLX_NO_SWF_VERSION_TEST);
 		defineInversion(FLX_NO_HEALTH, FLX_HEALTH);
+		defineInversion(FLX_TRACK_POOLS, FLX_NO_TRACK_POOLS);
 	}
 
 	static function defineHelperDefines()
@@ -237,6 +241,9 @@ class FlxDefines
 		// should always be defined as of 5.5.1 and, therefore, deprecated
 		define(FLX_DRAW_QUADS);
 		// #end
+		
+		if (defined(FLX_TRACK_POOLS) && !defined("debug"))
+			abort("Can only define FLX_TRACK_POOLS on debug mode", (macro null).pos);
 	}
 
 	static function defineInversion(userDefine:UserDefines, invertedDefine:HelperDefines)
diff --git a/flixel/util/FlxPool.hx b/flixel/util/FlxPool.hx
index 57609aa979..67ac7fae8d 100644
--- a/flixel/util/FlxPool.hx
+++ b/flixel/util/FlxPool.hx
@@ -1,6 +1,7 @@
 package flixel.util;
 
-import flixel.util.FlxDestroyUtil.IFlxDestroyable;
+import flixel.FlxG;
+import flixel.util.FlxDestroyUtil;
 import flixel.util.typeLimit.OneOfTwo;
 
 /**
@@ -47,6 +48,16 @@ class FlxPool<T:IFlxDestroyable> implements IFlxPool<T>
 	 * _count keeps track of the valid, accessible pool objects.
 	 */
 	var _count:Int = 0;
+	
+	#if FLX_TRACK_POOLS
+	/** Set to false before creating FlxGame to prevent logs */
+	public var autoLog = true;
+	
+	final _tracked = new Map<T, String>();
+	final _leakCount = new Map<String, Int>();
+	var _totalCreated = 0;
+	var _autoLogInitted = false;
+	#end
 
 	/**
 	 * Creates a pool of the specified type
@@ -61,11 +72,21 @@ class FlxPool<T:IFlxDestroyable> implements IFlxPool<T>
 
 	public function get():T
 	{
-		if (_count == 0)
+		final obj:T = if (_count == 0)
 		{
-			return _constructor();
+			#if FLX_TRACK_POOLS
+			_totalCreated++;
+			#end
+			_constructor();
 		}
-		return _pool[--_count];
+		else
+			_pool[--_count];
+		
+		#if FLX_TRACK_POOLS
+		trackGet(obj);
+		#end
+		
+		return obj;
 	}
 
 	public function put(obj:T):Void
@@ -91,6 +112,10 @@ class FlxPool<T:IFlxDestroyable> implements IFlxPool<T>
 	{
 		obj.destroy();
 		_pool[_count++] = obj;
+		
+		#if FLX_TRACK_POOLS
+		trackPut(obj);
+		#end
 	}
 
 	public function preAllocate(numObjects:Int):Void
@@ -113,6 +138,83 @@ class FlxPool<T:IFlxDestroyable> implements IFlxPool<T>
 	{
 		return _count;
 	}
+	
+	#if FLX_TRACK_POOLS
+	public function addLogs(?id:String)
+	{
+		if (id == null)
+		{
+			if (_pool.length == 0)
+				preAllocate(1);
+			
+			id = Type.getClassName(Type.getClass(_pool[0])).split(".").pop().split("Flx").pop();
+		}
+		
+		FlxG.watch.addFunction(id + "-pool", () -> '$length/$_totalCreated');
+		FlxG.watch.addFunction(id + "-top-leak", function()
+		{
+			var most = 0;
+			var topStack:String = null;
+			for (stack in _leakCount.keys())
+			{
+				final count = _leakCount[stack];
+				if (most < count)
+				{
+					most = count;
+					topStack = stack;
+				}
+			}
+			
+			return if (topStack == null)
+				null;
+			else
+				'${prettyStack(topStack)}: $most';
+		});
+	}
+	
+	function trackGet(obj:T)
+	{
+		final callStack = haxe.CallStack.callStack();
+		final stack = stackToString(callStack[3]);
+		if (stack == null)
+			return;
+		
+		if (autoLog && !_autoLogInitted && FlxG.signals != null)
+		{
+			_autoLogInitted = true;
+			FlxG.signals.postStateSwitch.add(()->addLogs());
+			if (FlxG.game != null && FlxG.state != null)
+				addLogs();
+		}
+		
+		_tracked[obj] = stack;
+		if (_leakCount.exists(stack) == false)
+			_leakCount[stack] = 0;
+		
+		_leakCount[stack]++;
+	}
+	
+	inline function trackPut(obj:T)
+	{
+		final stack = _tracked[obj];
+		_tracked.remove(obj);
+		_leakCount[stack]--;
+	}
+	
+	function stackToString(stack:haxe.CallStack.StackItem)
+	{
+		return switch (stack)
+		{
+			case FilePos(_, file, line, _): '$file[$line]';
+			default: null;
+		}
+	}
+	
+	inline function prettyStack(pos:String)
+	{
+		return topStack.split("/").pop().split(".hx").shift()
+	}
+	#end
 }
 
 interface IFlxPooled extends IFlxDestroyable
diff --git a/tests/coverage/Project.xml b/tests/coverage/Project.xml
index 87c127d82e..97a5198510 100644
--- a/tests/coverage/Project.xml
+++ b/tests/coverage/Project.xml
@@ -56,5 +56,6 @@
 	</section>
 	<section if="coverage2">
 		<haxedef name="debug" />
+		<haxedef name="FLX_TRACK_POOLS" />
 	</section>
 </project>