Add dominator APIs

2 files changed, 42 insertions, 9 deletions

diff --git a/python/basicblock.py b/python/basicblock.py
index ebacb311..9d48d3cb 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -134,6 +134,36 @@ class BasicBlock(object):
 		return core.BNBasicBlockHasUndeterminedOutgoingEdges(self.handle)
 
 	@property
+	def dominators(self):
+		"""List of dominators for this basic block (read-only)"""
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetBasicBlockDominators(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			result.append(BasicBlock(self.view, core.BNNewBasicBlockReference(blocks[i])))
+		core.BNFreeBasicBlockList(blocks, count.value)
+		return result
+
+	@property
+	def strict_dominators(self):
+		"""List of strict dominators for this basic block (read-only)"""
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetBasicBlockStrictDominators(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			result.append(BasicBlock(self.view, core.BNNewBasicBlockReference(blocks[i])))
+		core.BNFreeBasicBlockList(blocks, count.value)
+		return result
+
+	@property
+	def immediate_dominator(self):
+		"""Immediate dominator of this basic block (read-only)"""
+		result = core.BNGetBasicBlockImmediateDominator(self.handle)
+		if not result:
+			return None
+		return BasicBlock(self.view, result)
+
+	@property
 	def annotations(self):
 		"""List of automatic annotations for the start of this block (read-only)"""
 		return self.function.get_block_annotations(self.arch, self.start)
diff --git a/python/function.py b/python/function.py
index 4cdd6cc9..2f6788f6 100644
--- a/python/function.py
+++ b/python/function.py
@@ -847,16 +847,13 @@ class DisassemblyTextLine(object):
 
 
 class FunctionGraphEdge(object):
-	def __init__(self, branch_type, arch, target, points):
+	def __init__(self, branch_type, target, points):
 		self.type = BranchType(branch_type)
-		self.arch = arch
 		self.target = target
 		self.points = points
 
 	def __repr__(self):
-		if self.arch:
-			return "<%s: %s@%#x>" % (self.type.name, self.arch.name, self.target)
-		return "<%s: %#x>" % (self.type, self.target)
+		return "<%s: %s>" % (self.type.name, repr(self.target))
 
 
 class FunctionGraphBlock(object):
@@ -958,13 +955,19 @@ class FunctionGraphBlock(object):
 		for i in xrange(0, count.value):
 			branch_type = BranchType(edges[i].type)
 			target = edges[i].target
-			arch = None
-			if edges[i].arch is not None:
-				arch = architecture.Architecture(edges[i].arch)
+			if target:
+				func = core.BNGetBasicBlockFunction(target)
+				if func is None:
+					core.BNFreeBasicBlock(target)
+					target = None
+				else:
+					target = basicblock.BasicBlock(binaryview.BinaryView(handle = core.BNGetFunctionData(func)),
+						core.BNNewBasicBlockReference(target))
+					core.BNFreeFunction(func)
 			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, arch, target, points))
+			result.append(FunctionGraphEdge(branch_type, target, points))
 		core.BNFreeFunctionGraphBlockOutgoingEdgeList(edges, count.value)
 		return result