Add back edge property to block edges

3 files changed, 25 insertions, 6 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
 
diff --git a/python/examples/export_svg.py b/python/examples/export_svg.py
index 89bc41a1..95e2d62d 100755
--- a/python/examples/export_svg.py
+++ b/python/examples/export_svg.py
@@ -80,6 +80,10 @@ def render_svg(function):
 				fill: none;
 				stroke-width: 1px;
 			}
+			.back_edge {
+				fill: none;
+				stroke-width: 2px;
+			}
 			.UnconditionalBranch, .IndirectBranch {
 				stroke: rgb(128, 198, 233);
 				color: rgb(128, 198, 233);
@@ -197,7 +201,10 @@ def render_svg(function):
 				points += str(x * widthconst) + "," + str(y * heightconst) + " "
 			x, y = edge.points[-1]
 			points += str(x * widthconst) + "," + str(y * heightconst + 0) + " "
-			edges += '		<polyline class="edge {type}" points="{points}" marker-end="url(#arrow-{type})"/>\n'.format(type=BranchType(edge.type).name, points=points)
+			if edge.back_edge:
+				edges += '		<polyline class="back_edge {type}" points="{points}" marker-end="url(#arrow-{type})"/>\n'.format(type=BranchType(edge.type).name, points=points)
+			else:
+				edges += '		<polyline class="edge {type}" points="{points}" marker-end="url(#arrow-{type})"/>\n'.format(type=BranchType(edge.type).name, points=points)
 	output += ' ' + edges + '\n'
 	output += '	</g>\n'
 	output += '</svg></html>'
diff --git a/python/function.py b/python/function.py
index 2f6788f6..86c93bf7 100644
--- a/python/function.py
+++ b/python/function.py
@@ -847,14 +847,20 @@ class DisassemblyTextLine(object):
 
 
 class FunctionGraphEdge(object):
-	def __init__(self, branch_type, target, points):
+	def __init__(self, branch_type, source, target, points):
 		self.type = BranchType(branch_type)
+		self.source = source
 		self.target = target
 		self.points = points
 
 	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):
@@ -967,7 +973,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, target, points))
+			result.append(FunctionGraphEdge(branch_type, self, target, points))
 		core.BNFreeFunctionGraphBlockOutgoingEdgeList(edges, count.value)
 		return result