4 files changed, 92 insertions, 0 deletions

diff --git a/basicblock.cpp b/basicblock.cpp
index d6423fa0..d8ae0f65 100644
--- a/basicblock.cpp
+++ b/basicblock.cpp
@@ -195,6 +195,54 @@ Ref<BasicBlock> BasicBlock::GetImmediateDominator() const
 }
 
 
+set<Ref<BasicBlock>> BasicBlock::GetDominatorTreeChildren() const
+{
+	size_t count;
+	BNBasicBlock** blocks = BNGetBasicBlockDominatorTreeChildren(m_object, &count);
+
+	set<Ref<BasicBlock>> result;
+	for (size_t i = 0; i < count; i++)
+		result.insert(new BasicBlock(BNNewBasicBlockReference(blocks[i])));
+
+	BNFreeBasicBlockList(blocks, count);
+	return result;
+}
+
+
+set<Ref<BasicBlock>> BasicBlock::GetDominanceFrontier() const
+{
+	size_t count;
+	BNBasicBlock** blocks = BNGetBasicBlockDominanceFrontier(m_object, &count);
+
+	set<Ref<BasicBlock>> result;
+	for (size_t i = 0; i < count; i++)
+		result.insert(new BasicBlock(BNNewBasicBlockReference(blocks[i])));
+
+	BNFreeBasicBlockList(blocks, count);
+	return result;
+}
+
+
+set<Ref<BasicBlock>> BasicBlock::GetIteratedDominanceFrontier(const set<Ref<BasicBlock>>& blocks)
+{
+	BNBasicBlock** blockSet = new BNBasicBlock*[blocks.size()];
+	size_t i = 0;
+	for (auto& j : blocks)
+		blockSet[i++] = j->GetObject();
+
+	size_t count;
+	BNBasicBlock** resultBlocks = BNGetBasicBlockIteratedDominanceFrontier(blockSet, blocks.size(), &count);
+	delete[] blockSet;
+
+	set<Ref<BasicBlock>> result;
+	for (size_t i = 0; i < count; i++)
+		result.insert(new BasicBlock(BNNewBasicBlockReference(resultBlocks[i])));
+
+	BNFreeBasicBlockList(resultBlocks, count);
+	return result;
+}
+
+
 void BasicBlock::MarkRecentUse()
 {
 	BNMarkBasicBlockAsRecentlyUsed(m_object);
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 6a23fd13..6b69c95f 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1784,6 +1784,9 @@ namespace BinaryNinja
 		std::set<Ref<BasicBlock>> GetDominators() const;
 		std::set<Ref<BasicBlock>> GetStrictDominators() const;
 		Ref<BasicBlock> GetImmediateDominator() const;
+		std::set<Ref<BasicBlock>> GetDominatorTreeChildren() const;
+		std::set<Ref<BasicBlock>> GetDominanceFrontier() const;
+		static std::set<Ref<BasicBlock>> GetIteratedDominanceFrontier(const std::set<Ref<BasicBlock>>& blocks);
 
 		void MarkRecentUse();
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 239d1002..d1c7f4d2 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1731,6 +1731,10 @@ extern "C"
 	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlockDominators(BNBasicBlock* block, size_t* count);
 	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlockStrictDominators(BNBasicBlock* block, size_t* count);
 	BINARYNINJACOREAPI BNBasicBlock* BNGetBasicBlockImmediateDominator(BNBasicBlock* block);
+	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlockDominatorTreeChildren(BNBasicBlock* block, size_t* count);
+	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlockDominanceFrontier(BNBasicBlock* block, size_t* count);
+	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlockIteratedDominanceFrontier(BNBasicBlock** blocks,
+		size_t incomingCount, size_t* outputCount);
 
 	BINARYNINJACOREAPI BNDisassemblyTextLine* BNGetBasicBlockDisassemblyText(BNBasicBlock* block,
 		BNDisassemblySettings* settings, size_t* count);
diff --git a/python/basicblock.py b/python/basicblock.py
index e332cc3c..9f256aa7 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -170,6 +170,28 @@ class BasicBlock(object):
 		return BasicBlock(self.view, result)
 
 	@property
+	def dominator_tree_children(self):
+		"""List of child blocks in the dominator tree for this basic block (read-only)"""
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetBasicBlockDominatorTreeChildren(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 dominance_frontier(self):
+		"""Dominance frontier for this basic block (read-only)"""
+		count = ctypes.c_ulonglong()
+		blocks = core.BNGetBasicBlockDominanceFrontier(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 annotations(self):
 		"""List of automatic annotations for the start of this block (read-only)"""
 		return self.function.get_block_annotations(self.arch, self.start)
@@ -208,6 +230,21 @@ class BasicBlock(object):
 	def highlight(self, value):
 		self.set_user_highlight(value)
 
+	@classmethod
+	def get_iterated_dominance_frontier(self, blocks):
+		if len(blocks) == 0:
+			return []
+		block_set = (ctypes.POINTER(core.BNBasicBlock) * len(blocks))()
+		for i in xrange(len(blocks)):
+			block_set[i] = blocks[i].handle
+		count = ctypes.c_ulonglong()
+		out_blocks = core.BNGetBasicBlockIteratedDominanceFrontier(block_set, len(blocks), count)
+		result = []
+		for i in xrange(0, count.value):
+			result.append(BasicBlock(blocks[0].view, core.BNNewBasicBlockReference(out_blocks[i])))
+		core.BNFreeBasicBlockList(out_blocks, count.value)
+		return result
+
 	def __setattr__(self, name, value):
 		try:
 			object.__setattr__(self, name, value)