Compute back edges in the core

6 files changed, 20 insertions, 16 deletions

diff --git a/basicblock.cpp b/basicblock.cpp
index 7eb37c57..c2a2bddf 100644
--- a/basicblock.cpp
+++ b/basicblock.cpp
@@ -125,6 +125,7 @@ vector<BasicBlockEdge> BasicBlock::GetOutgoingEdges() const
 		BasicBlockEdge edge;
 		edge.type = array[i].type;
 		edge.target = array[i].target ? new BasicBlock(BNNewBasicBlockReference(array[i].target)) : nullptr;
+		edge.backEdge = array[i].backEdge;
 		result.push_back(edge);
 	}
 
@@ -144,6 +145,7 @@ vector<BasicBlockEdge> BasicBlock::GetIncomingEdges() const
 		BasicBlockEdge edge;
 		edge.type = array[i].type;
 		edge.target = array[i].target ? new BasicBlock(BNNewBasicBlockReference(array[i].target)) : nullptr;
+		edge.backEdge = array[i].backEdge;
 		result.push_back(edge);
 	}
 
@@ -397,5 +399,10 @@ void BasicBlock::SetUserBasicBlockHighlight(uint8_t r, uint8_t g, uint8_t b, uin
 
 bool BasicBlock::IsBackEdge(BasicBlock* source, BasicBlock* target)
 {
-	return source->GetDominators().count(target) != 0;
+	for (auto& i : source->GetOutgoingEdges())
+	{
+		if (i.target->GetObject() == target->GetObject())
+			return i.backEdge;
+	}
+	return false;
 }
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 5ea6d909..e5ae77f5 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1784,6 +1784,7 @@ namespace BinaryNinja
 	{
 		BNBranchType type;
 		Ref<BasicBlock> target;
+		bool backEdge;
 	};
 
 	class BasicBlock: public CoreRefCountObject<BNBasicBlock, BNNewBasicBlockReference, BNFreeBasicBlock>
@@ -2041,6 +2042,7 @@ namespace BinaryNinja
 		BNBranchType type;
 		Ref<BasicBlock> target;
 		std::vector<BNPoint> points;
+		bool backEdge;
 	};
 
 	class FunctionGraphBlock: public CoreRefCountObject<BNFunctionGraphBlock,
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 8880747c..d053eb2b 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1021,6 +1021,7 @@ extern "C"
 	{
 		BNBranchType type;
 		BNBasicBlock* target;
+		bool backEdge;
 	};
 
 	struct BNPoint
@@ -1035,6 +1036,7 @@ extern "C"
 		BNBasicBlock* target;
 		BNPoint* points;
 		size_t pointCount;
+		bool backEdge;
 	};
 
 	struct BNDisassemblyTextLine
diff --git a/functiongraphblock.cpp b/functiongraphblock.cpp
index 99ce2e08..938fb635 100644
--- a/functiongraphblock.cpp
+++ b/functiongraphblock.cpp
@@ -130,6 +130,7 @@ const vector<FunctionGraphEdge>& FunctionGraphBlock::GetOutgoingEdges()
 		edge.type = edges[i].type;
 		edge.target = edges[i].target ? new BasicBlock(BNNewBasicBlockReference(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);
 	}
 
diff --git a/python/basicblock.py b/python/basicblock.py
index 3dc5b050..7858cc60 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -29,10 +29,11 @@ import function
 
 
 class BasicBlockEdge(object):
-	def __init__(self, branch_type, source, target):
+	def __init__(self, branch_type, source, target, back_edge):
 		self.type = branch_type
 		self.source = source
 		self.target = target
+		self.back_edge = back_edge
 
 	def __repr__(self):
 		if self.type == BranchType.UnresolvedBranch:
@@ -42,11 +43,6 @@ class BasicBlockEdge(object):
 		else:
 			return "<%s: %#x>" % (BranchType(self.type).name, self.target.start)
 
-	@property
-	def back_edge(self):
-		"""Whether the edge is a back edge (end of a loop)"""
-		return self.target in self.source.dominators
-
 
 class BasicBlock(object):
 	def __init__(self, view, handle):
@@ -114,7 +110,7 @@ class BasicBlock(object):
 				target = BasicBlock(self.view, core.BNNewBasicBlockReference(edges[i].target))
 			else:
 				target = None
-			result.append(BasicBlockEdge(branch_type, self, target))
+			result.append(BasicBlockEdge(branch_type, self, target, edges[i].backEdge))
 		core.BNFreeBasicBlockEdgeList(edges, count.value)
 		return result
 
@@ -130,7 +126,7 @@ class BasicBlock(object):
 				target = BasicBlock(self.view, core.BNNewBasicBlockReference(edges[i].target))
 			else:
 				target = None
-			result.append(BasicBlockEdge(branch_type, self, target))
+			result.append(BasicBlockEdge(branch_type, self, target, edges[i].backEdge))
 		core.BNFreeBasicBlockEdgeList(edges, count.value)
 		return result
 
diff --git a/python/function.py b/python/function.py
index 34511cfa..8beebe66 100644
--- a/python/function.py
+++ b/python/function.py
@@ -948,20 +948,16 @@ class DisassemblyTextLine(object):
 
 
 class FunctionGraphEdge(object):
-	def __init__(self, branch_type, source, target, points):
+	def __init__(self, branch_type, source, target, points, back_edge):
 		self.type = BranchType(branch_type)
 		self.source = source
 		self.target = target
 		self.points = points
+		self.back_edge = back_edge
 
 	def __repr__(self):
 		return "<%s: %s>" % (self.type.name, repr(self.target))
 
-	@property
-	def back_edge(self):
-		"""Whether the edge is a back edge (end of a loop)"""
-		return self.target in self.source.basic_block.dominators
-
 
 class FunctionGraphBlock(object):
 	def __init__(self, handle):
@@ -1074,7 +1070,7 @@ class FunctionGraphBlock(object):
 			points = []
 			for j in xrange(0, edges[i].pointCount):
 				points.append((edges[i].points[j].x, edges[i].points[j].y))
-			result.append(FunctionGraphEdge(branch_type, self, target, points))
+			result.append(FunctionGraphEdge(branch_type, self, target, points, edges[i].backEdge))
 		core.BNFreeFunctionGraphBlockOutgoingEdgeList(edges, count.value)
 		return result