Add back edge property to block edges

1 files changed, 9 insertions, 3 deletions

diff --git a/python/basicblock.py b/python/basicblock.py
index 9d48d3cb..e332cc3c 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -29,8 +29,9 @@ import function
 
 
 class BasicBlockEdge(object):
-	def __init__(self, branch_type, target):
+	def __init__(self, branch_type, source, target):
 		self.type = branch_type
+		self.source = source
 		self.target = target
 
 	def __repr__(self):
@@ -41,6 +42,11 @@ 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):
@@ -108,7 +114,7 @@ class BasicBlock(object):
 				target = BasicBlock(self.view, core.BNNewBasicBlockReference(edges[i].target))
 			else:
 				target = None
-			result.append(BasicBlockEdge(branch_type, target))
+			result.append(BasicBlockEdge(branch_type, self, target))
 		core.BNFreeBasicBlockEdgeList(edges, count.value)
 		return result
 
@@ -124,7 +130,7 @@ class BasicBlock(object):
 				target = BasicBlock(self.view, core.BNNewBasicBlockReference(edges[i].target))
 			else:
 				target = None
-			result.append(BasicBlockEdge(branch_type, target))
+			result.append(BasicBlockEdge(branch_type, self, target))
 		core.BNFreeBasicBlockEdgeList(edges, count.value)
 		return result