From 7887727f48ab8078a1ac855bb4b0766af7e2d619 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty@vector35.com>
Date: Sat, 27 Apr 2019 20:41:56 -0400
Subject: Add debug mode for finding refcount leaks

---
 binaryninjaapi.h  | 32 ++++++++++++++++++++++++++++++++
 binaryninjacore.h |  3 +++
 2 files changed, 35 insertions(+)

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 1b0232fc..b6d71165 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -43,6 +43,8 @@
 #define NOEXCEPT noexcept
 #endif
 
+//#define BN_REF_COUNT_DEBUG  // Mac OS X only, prints stack trace of leaked references
+
 
 namespace BinaryNinja
 {
@@ -199,6 +201,9 @@ namespace BinaryNinja
 	class Ref
 	{
 		T* m_obj;
+#ifdef BN_REF_COUNT_DEBUG
+		void* m_assignmentTrace = nullptr;
+#endif
 
 	public:
 		Ref<T>(): m_obj(NULL)
@@ -208,23 +213,44 @@ namespace BinaryNinja
 		Ref<T>(T* obj): m_obj(obj)
 		{
 			if (m_obj)
+			{
 				m_obj->AddRef();
+#ifdef BN_REF_COUNT_DEBUG
+				m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
+			}
 		}
 
 		Ref<T>(const Ref<T>& obj): m_obj(obj.m_obj)
 		{
 			if (m_obj)
+			{
 				m_obj->AddRef();
+#ifdef BN_REF_COUNT_DEBUG
+				m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
+			}
 		}
 
 		~Ref<T>()
 		{
 			if (m_obj)
+			{
 				m_obj->Release();
+#ifdef BN_REF_COUNT_DEBUG
+				BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+#endif
+			}
 		}
 
 		Ref<T>& operator=(const Ref<T>& obj)
 		{
+#ifdef BN_REF_COUNT_DEBUG
+			if (m_obj)
+				BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+			if (obj.m_obj)
+				m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
 			T* oldObj = m_obj;
 			m_obj = obj.m_obj;
 			if (m_obj)
@@ -236,6 +262,12 @@ namespace BinaryNinja
 
 		Ref<T>& operator=(T* obj)
 		{
+#ifdef BN_REF_COUNT_DEBUG
+			if (m_obj)
+				BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+			if (obj)
+				m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
 			T* oldObj = m_obj;
 			m_obj = obj;
 			if (m_obj)
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 89b61202..383e6fb2 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3921,6 +3921,9 @@ extern "C"
 	BINARYNINJACOREAPI bool BNParseExpression(BNBinaryView* view, const char* expression, uint64_t* offset, uint64_t here, char** errorString);
 	BINARYNINJACOREAPI void BNFreeParseError(char* errorString);
 
+	BINARYNINJACOREAPI void* BNRegisterObjectRefDebugTrace(const char* typeName);
+	BINARYNINJACOREAPI void BNUnregisterObjectRefDebugTrace(const char* typeName, void* trace);
+
 #ifdef __cplusplus
 }
 #endif
-- 
cgit v1.3.1