Add incoming edges for flow graph nodes

4 files changed, 54 insertions, 5 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 6c7b05b8..59f0313b 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2718,8 +2718,8 @@ namespace BinaryNinja
 		BNNewFlowGraphNodeReference, BNFreeFlowGraphNode>
 	{
 		std::vector<DisassemblyTextLine> m_cachedLines;
-		std::vector<FlowGraphEdge> m_cachedEdges;
-		bool m_cachedLinesValid, m_cachedEdgesValid;
+		std::vector<FlowGraphEdge> m_cachedEdges, m_cachedIncomingEdges;
+		bool m_cachedLinesValid, m_cachedEdgesValid, m_cachedIncomingEdgesValid;
 
 	public:
 		FlowGraphNode(FlowGraph* graph);
@@ -2736,6 +2736,7 @@ namespace BinaryNinja
 		const std::vector<DisassemblyTextLine>& GetLines();
 		void SetLines(const std::vector<DisassemblyTextLine>& lines);
 		const std::vector<FlowGraphEdge>& GetOutgoingEdges();
+		const std::vector<FlowGraphEdge>& GetIncomingEdges();
 		void AddOutgoingEdge(BNBranchType type, FlowGraphNode* target);
 
 		BNHighlightColor GetHighlight() const;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index c1d993e5..0ba3fded 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2933,7 +2933,8 @@ extern "C"
 	BINARYNINJACOREAPI BNDisassemblyTextLine* BNGetFlowGraphNodeLines(BNFlowGraphNode* node, size_t* count);
 	BINARYNINJACOREAPI void BNSetFlowGraphNodeLines(BNFlowGraphNode* node, BNDisassemblyTextLine* lines, size_t count);
 	BINARYNINJACOREAPI BNFlowGraphEdge* BNGetFlowGraphNodeOutgoingEdges(BNFlowGraphNode* node, size_t* count);
-	BINARYNINJACOREAPI void BNFreeFlowGraphNodeOutgoingEdgeList(BNFlowGraphEdge* edges, size_t count);
+	BINARYNINJACOREAPI BNFlowGraphEdge* BNGetFlowGraphNodeIncomingEdges(BNFlowGraphNode* node, size_t* count);
+	BINARYNINJACOREAPI void BNFreeFlowGraphNodeEdgeList(BNFlowGraphEdge* edges, size_t count);
 	BINARYNINJACOREAPI void BNAddFlowGraphNodeOutgoingEdge(BNFlowGraphNode* node, BNBranchType type, BNFlowGraphNode* target);
 
 	BINARYNINJACOREAPI BNHighlightColor BNGetFlowGraphNodeHighlight(BNFlowGraphNode* node);
diff --git a/flowgraphnode.cpp b/flowgraphnode.cpp
index 0101d716..37128cd8 100644
--- a/flowgraphnode.cpp
+++ b/flowgraphnode.cpp
@@ -29,6 +29,7 @@ FlowGraphNode::FlowGraphNode(FlowGraph* graph)
 	m_object = BNCreateFlowGraphNode(graph->GetObject());
 	m_cachedLinesValid = false;
 	m_cachedEdgesValid = false;
+	m_cachedIncomingEdgesValid = false;
 }
 
 
@@ -37,6 +38,7 @@ FlowGraphNode::FlowGraphNode(BNFlowGraphNode* node)
 	m_object = node;
 	m_cachedLinesValid = false;
 	m_cachedEdgesValid = false;
+	m_cachedIncomingEdgesValid = false;
 }
 
 
@@ -154,13 +156,40 @@ const vector<FlowGraphEdge>& FlowGraphNode::GetOutgoingEdges()
 		result.push_back(edge);
 	}
 
-	BNFreeFlowGraphNodeOutgoingEdgeList(edges, count);
+	BNFreeFlowGraphNodeEdgeList(edges, count);
 	m_cachedEdges = result;
 	m_cachedEdgesValid = true;
 	return m_cachedEdges;
 }
 
 
+const vector<FlowGraphEdge>& FlowGraphNode::GetIncomingEdges()
+{
+	if (m_cachedIncomingEdgesValid)
+		return m_cachedIncomingEdges;
+
+	size_t count;
+	BNFlowGraphEdge* edges = BNGetFlowGraphNodeIncomingEdges(m_object, &count);
+
+	vector<FlowGraphEdge> result;
+	result.reserve(count);
+	for (size_t i = 0; i < count; i++)
+	{
+		FlowGraphEdge edge;
+		edge.type = edges[i].type;
+		edge.target = edges[i].target ? new FlowGraphNode(BNNewFlowGraphNodeReference(edges[i].target)) : nullptr;
+		edge.points.insert(edge.points.begin(), &edges[i].points[0], &edges[i].points[edges[i].pointCount]);
+		edge.backEdge = edges[i].backEdge;
+		result.push_back(edge);
+	}
+
+	BNFreeFlowGraphNodeEdgeList(edges, count);
+	m_cachedIncomingEdges = result;
+	m_cachedIncomingEdgesValid = true;
+	return m_cachedIncomingEdges;
+}
+
+
 void FlowGraphNode::AddOutgoingEdge(BNBranchType type, FlowGraphNode* target)
 {
 	BNAddFlowGraphNodeOutgoingEdge(m_object, type, target->GetObject());
diff --git a/python/flowgraph.py b/python/flowgraph.py
index ce92e47b..3054ff35 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -193,7 +193,25 @@ class FlowGraphNode(object):
 			for j in range(0, edges[i].pointCount):
 				points.append((edges[i].points[j].x, edges[i].points[j].y))
 			result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge))
-		core.BNFreeFlowGraphNodeOutgoingEdgeList(edges, count.value)
+		core.BNFreeFlowGraphNodeEdgeList(edges, count.value)
+		return result
+
+	@property
+	def incoming_edges(self):
+		"""Flow graph block list of incoming edges (read-only)"""
+		count = ctypes.c_ulonglong()
+		edges = core.BNGetFlowGraphNodeIncomingEdges(self.handle, count)
+		result = []
+		for i in range(0, count.value):
+			branch_type = BranchType(edges[i].type)
+			target = edges[i].target
+			if target:
+				target = FlowGraphNode(self.graph, core.BNNewFlowGraphNodeReference(target))
+			points = []
+			for j in range(0, edges[i].pointCount):
+				points.append((edges[i].points[j].x, edges[i].points[j].y))
+			result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge))
+		core.BNFreeFlowGraphNodeEdgeList(edges, count.value)
 		return result
 
 	@property