1 files changed, 121 insertions, 4 deletions

diff --git a/basicblock.cpp b/basicblock.cpp
index 93a279a1..7eb37c57 100644
--- a/basicblock.cpp
+++ b/basicblock.cpp
@@ -1,4 +1,4 @@
-// Copyright (c) 2015-2016 Vector 35 LLC
+// Copyright (c) 2015-2017 Vector 35 LLC
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy
 // of this software and associated documentation files (the "Software"), to
@@ -108,6 +108,12 @@ uint64_t BasicBlock::GetLength() const
 }
 
 
+size_t BasicBlock::GetIndex() const
+{
+	return BNGetBasicBlockIndex(m_object);
+}
+
+
 vector<BasicBlockEdge> BasicBlock::GetOutgoingEdges() const
 {
 	size_t count;
@@ -118,12 +124,30 @@ vector<BasicBlockEdge> BasicBlock::GetOutgoingEdges() const
 	{
 		BasicBlockEdge edge;
 		edge.type = array[i].type;
-		edge.target = array[i].target;
-		edge.arch = array[i].arch ? new CoreArchitecture(array[i].arch) : nullptr;
+		edge.target = array[i].target ? new BasicBlock(BNNewBasicBlockReference(array[i].target)) : nullptr;
+		result.push_back(edge);
+	}
+
+	BNFreeBasicBlockEdgeList(array, count);
+	return result;
+}
+
+
+vector<BasicBlockEdge> BasicBlock::GetIncomingEdges() const
+{
+	size_t count;
+	BNBasicBlockEdge* array = BNGetBasicBlockIncomingEdges(m_object, &count);
+
+	vector<BasicBlockEdge> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		BasicBlockEdge edge;
+		edge.type = array[i].type;
+		edge.target = array[i].target ? new BasicBlock(BNNewBasicBlockReference(array[i].target)) : nullptr;
 		result.push_back(edge);
 	}
 
-	BNFreeBasicBlockOutgoingEdgeList(array);
+	BNFreeBasicBlockEdgeList(array, count);
 	return result;
 }
 
@@ -134,6 +158,91 @@ bool BasicBlock::HasUndeterminedOutgoingEdges() const
 }
 
 
+set<Ref<BasicBlock>> BasicBlock::GetDominators() const
+{
+	size_t count;
+	BNBasicBlock** blocks = BNGetBasicBlockDominators(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::GetStrictDominators() const
+{
+	size_t count;
+	BNBasicBlock** blocks = BNGetBasicBlockStrictDominators(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;
+}
+
+
+Ref<BasicBlock> BasicBlock::GetImmediateDominator() const
+{
+	BNBasicBlock* result = BNGetBasicBlockImmediateDominator(m_object);
+	if (!result)
+		return nullptr;
+	return new BasicBlock(result);
+}
+
+
+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);
@@ -164,6 +273,8 @@ vector<DisassemblyTextLine> BasicBlock::GetDisassemblyText(DisassemblySettings*
 			token.value = lines[i].tokens[j].value;
 			token.size = lines[i].tokens[j].size;
 			token.operand = lines[i].tokens[j].operand;
+			token.context = lines[i].tokens[j].context;
+			token.address = lines[i].tokens[j].address;
 			line.tokens.push_back(token);
 		}
 		result.push_back(line);
@@ -282,3 +393,9 @@ void BasicBlock::SetUserBasicBlockHighlight(uint8_t r, uint8_t g, uint8_t b, uin
 	hc.alpha = alpha;
 	SetUserBasicBlockHighlight(hc);
 }
+
+
+bool BasicBlock::IsBackEdge(BasicBlock* source, BasicBlock* target)
+{
+	return source->GetDominators().count(target) != 0;
+}