Add APIs for subclassing flow graphs

4 files changed, 134 insertions, 9 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 3582b265..1133625e 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -54,6 +54,7 @@ namespace BinaryNinja
 		virtual ~RefCountObject() {}
 
 		RefCountObject* GetObject() { return this; }
+		static RefCountObject* GetObject(RefCountObject* obj) { return obj; }
 
 		void AddRef()
 		{
@@ -107,6 +108,13 @@ namespace BinaryNinja
 
 		T* GetObject() const { return m_object; }
 
+		static T* GetObject(CoreRefCountObject* obj)
+		{
+			if (!obj)
+				return nullptr;
+			return obj->GetObject();
+		}
+
 		void AddRef()
 		{
 			if (m_object && (m_refs != 0))
@@ -164,6 +172,13 @@ namespace BinaryNinja
 
 		T* GetObject() const { return m_object; }
 
+		static T* GetObject(StaticCoreRefCountObject* obj)
+		{
+			if (!obj)
+				return nullptr;
+			return obj->GetObject();
+		}
+
 		void AddRef()
 		{
 			AddRefInternal();
@@ -252,32 +267,32 @@ namespace BinaryNinja
 
 		bool operator==(const T* obj) const
 		{
-			return m_obj->GetObject() == obj->GetObject();
+			return T::GetObject(m_obj) == T::GetObject(obj);
 		}
 
 		bool operator==(const Ref<T>& obj) const
 		{
-			return m_obj->GetObject() == obj.m_obj->GetObject();
+			return T::GetObject(m_obj) == T::GetObject(obj.m_obj);
 		}
 
 		bool operator!=(const T* obj) const
 		{
-			return m_obj->GetObject() != obj->GetObject();
+			return T::GetObject(m_obj) != T::GetObject(obj);
 		}
 
 		bool operator!=(const Ref<T>& obj) const
 		{
-			return m_obj->GetObject() != obj.m_obj->GetObject();
+			return T::GetObject(m_obj) != T::GetObject(obj.m_obj);
 		}
 
 		bool operator<(const T* obj) const
 		{
-			return m_obj->GetObject() < obj->GetObject();
+			return T::GetObject(m_obj) < T::GetObject(obj);
 		}
 
 		bool operator<(const Ref<T>& obj) const
 		{
-			return m_obj->GetObject() < obj.m_obj->GetObject();
+			return T::GetObject(m_obj) < T::GetObject(obj.m_obj);
 		}
 
 		T* GetPtr() const
@@ -2567,7 +2582,18 @@ namespace BinaryNinja
 
 		static void CompleteCallback(void* ctxt);
 
+		static void PrepareForLayoutCallback(void* ctxt);
+		static void PopulateNodesCallback(void* ctxt);
+		static void CompleteLayoutCallback(void* ctxt);
+
+	protected:
+		void FinishPrepareForLayout();
+		virtual void PrepareForLayout();
+		virtual void PopulateNodes();
+		virtual void CompleteLayout();
+
 	public:
+		FlowGraph();
 		FlowGraph(BNFlowGraph* graph);
 		~FlowGraph();
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index c057b378..a32d8364 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1772,6 +1772,14 @@ extern "C"
 		FlowGraphReportType
 	};
 
+	struct BNCustomFlowGraph
+	{
+		void* context;
+		void (*prepareForLayout)(void* ctxt);
+		void (*populateNodes)(void* ctxt);
+		void (*completeLayout)(void* ctxt);
+	};
+
 	BINARYNINJACOREAPI char* BNAllocString(const char* contents);
 	BINARYNINJACOREAPI void BNFreeString(char* str);
 	BINARYNINJACOREAPI char** BNAllocStringList(const char** contents, size_t size);
@@ -2566,6 +2574,7 @@ extern "C"
 	BINARYNINJACOREAPI BNFlowGraph* BNCreateFlowGraph();
 	BINARYNINJACOREAPI BNFlowGraph* BNCreateFunctionGraph(BNFunction* func, BNFunctionGraphType type,
 		BNDisassemblySettings* settings);
+	BINARYNINJACOREAPI BNFlowGraph* BNCreateCustomFlowGraph(BNCustomFlowGraph* callbacks);
 	BINARYNINJACOREAPI BNFlowGraph* BNNewFlowGraphReference(BNFlowGraph* graph);
 	BINARYNINJACOREAPI void BNFreeFlowGraph(BNFlowGraph* graph);
 	BINARYNINJACOREAPI BNFunction* BNGetFunctionForFlowGraph(BNFlowGraph* graph);
@@ -2618,6 +2627,8 @@ extern "C"
 	BINARYNINJACOREAPI BNHighlightColor BNGetFlowGraphNodeHighlight(BNFlowGraphNode* node);
 	BINARYNINJACOREAPI void BNSetFlowGraphNodeHighlight(BNFlowGraphNode* node, BNHighlightColor color);
 
+	BINARYNINJACOREAPI void BNFinishPrepareForLayout(BNFlowGraph* graph);
+
 	// Symbols
 	BINARYNINJACOREAPI BNSymbol* BNCreateSymbol(BNSymbolType type, const char* shortName, const char* fullName,
 	                                            const char* rawName, uint64_t addr);
diff --git a/flowgraph.cpp b/flowgraph.cpp
index e557f540..4bacee81 100644
--- a/flowgraph.cpp
+++ b/flowgraph.cpp
@@ -24,6 +24,17 @@ using namespace BinaryNinja;
 using namespace std;
 
 
+FlowGraph::FlowGraph()
+{
+	BNCustomFlowGraph callbacks;
+	callbacks.context = this;
+	callbacks.prepareForLayout = PrepareForLayoutCallback;
+	callbacks.populateNodes = PopulateNodesCallback;
+	callbacks.completeLayout = CompleteLayoutCallback;
+	m_graph = BNCreateCustomFlowGraph(&callbacks);
+}
+
+
 FlowGraph::FlowGraph(BNFlowGraph* graph): m_graph(graph)
 {
 }
@@ -46,6 +57,49 @@ void FlowGraph::CompleteCallback(void* ctxt)
 }
 
 
+void FlowGraph::PrepareForLayoutCallback(void* ctxt)
+{
+	FlowGraph* graph = (FlowGraph*)ctxt;
+	graph->PrepareForLayout();
+}
+
+
+void FlowGraph::PopulateNodesCallback(void* ctxt)
+{
+	FlowGraph* graph = (FlowGraph*)ctxt;
+	graph->PopulateNodes();
+}
+
+
+void FlowGraph::CompleteLayoutCallback(void* ctxt)
+{
+	FlowGraph* graph = (FlowGraph*)ctxt;
+	graph->CompleteLayout();
+}
+
+
+void FlowGraph::FinishPrepareForLayout()
+{
+	BNFinishPrepareForLayout(m_graph);
+}
+
+
+void FlowGraph::PrepareForLayout()
+{
+	FinishPrepareForLayout();
+}
+
+
+void FlowGraph::PopulateNodes()
+{
+}
+
+
+void FlowGraph::CompleteLayout()
+{
+}
+
+
 Ref<Function> FlowGraph::GetFunction() const
 {
 	BNFunction* func = BNGetFunctionForFlowGraph(m_graph);
diff --git a/python/flowgraph.py b/python/flowgraph.py
index 1568bc77..25fdc430 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -24,13 +24,12 @@ import traceback
 
 # Binary Ninja components
 import _binaryninjacore as core
-from enums import (BranchType, InstructionTextTokenType, HighlightColorStyle, HighlightStandardColor)
+from enums import (BranchType, InstructionTextTokenType, HighlightStandardColor)
 import function
 import binaryview
 import lowlevelil
 import mediumlevelil
 import basicblock
-import architecture
 import log
 import interaction
 import highlight
@@ -264,7 +263,12 @@ class FlowGraphNode(object):
 class FlowGraph(object):
 	def __init__(self, handle = None):
 		if handle is None:
-			handle = core.BNCreateFlowGraph()
+			self._ext_cb = core.BNCustomFlowGraph()
+			self._ext_cb.context = 0
+			self._ext_cb.prepareForLayout = self._ext_cb.prepareForLayout.__class__(self._prepare_for_layout)
+			self._ext_cb.populateNodes = self._ext_cb.populateNodes.__class__(self._populate_nodes)
+			self._ext_cb.completeLayout = self._ext_cb.completeLayout.__class__(self._complete_layout)
+			handle = core.BNCreateCustomFlowGraph(self._ext_cb)
 		self.handle = handle
 		self._on_complete = None
 		self._cb = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self._complete)
@@ -283,6 +287,36 @@ class FlowGraph(object):
 			return True
 		return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents)
 
+	def _prepare_for_layout(self, ctxt):
+		try:
+			self.prepare_for_layout()
+		except:
+			log.log_error(traceback.format_exc())
+
+	def _populate_nodes(self, ctxt):
+		try:
+			self.populate_nodes()
+		except:
+			log.log_error(traceback.format_exc())
+
+	def _complete_layout(self, ctxt):
+		try:
+			self.complete_layout()
+		except:
+			log.log_error(traceback.format_exc())
+
+	def finish_prepare_for_layout(self):
+		core.BNFinishPrepareForLayout(self.handle)
+
+	def prepare_for_layout(self):
+		self.finish_prepare_for_layout()
+
+	def populate_nodes(self):
+		pass
+
+	def complete_layout(self):
+		pass
+
 	@property
 	def function(self):
 		"""Function for a flow graph"""