Initial code for generating function graphs

5 files changed, 329 insertions, 1 deletions

diff --git a/api.pro b/api.pro
index 8dfd78bd..0fe395f2 100644
--- a/api.pro
+++ b/api.pro
@@ -22,7 +22,9 @@ SOURCES += \
     transform.cpp \
 	architecture.cpp \
 	basicblock.cpp \
-	function.cpp
+	function.cpp \
+	functiongraph.cpp \
+	functiongraphblock.cpp
 
 HEADERS += binaryninjaapi.h
 unix {
diff --git a/architecture.cpp b/architecture.cpp
index 60fa4efa..0d010d0e 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -20,6 +20,12 @@ void InstructionInfo::AddBranch(BNBranchType type, uint64_t target)
 }
 
 
+InstructionTextToken::InstructionTextToken(BNInstructionTextTokenType t, const std::string& txt, uint64_t val) :
+	type(t), text(txt), value(val)
+{
+}
+
+
 Architecture::Architecture(BNArchitecture* arch): m_arch(arch)
 {
 }
@@ -42,11 +48,48 @@ bool Architecture::GetInstructionInfoCallback(void* ctxt, BNBinaryView* data, ui
 }
 
 
+bool Architecture::GetInstructionTextCallback(void* ctxt, BNBinaryView* data, uint64_t addr,
+                                              BNInstructionTextToken** result, size_t* count)
+{
+	Architecture* arch = (Architecture*)ctxt;
+	Ref<BinaryView> view = new BinaryView(BNNewViewReference(data));
+
+	vector<InstructionTextToken> tokens;
+	bool ok = arch->GetInstructionText(view, addr, tokens);
+	if (!ok)
+	{
+		*result = nullptr;
+		*count = 0;
+		return false;
+	}
+
+	*count = tokens.size();
+	*result = new BNInstructionTextToken[tokens.size()];
+	for (size_t i = 0; i < tokens.size(); i++)
+	{
+		(*result)[i].type = tokens[i].type;
+		(*result)[i].text = BNAllocString(tokens[i].text.c_str());
+		(*result)[i].value = tokens[i].value;
+	}
+	return true;
+}
+
+
+void Architecture::FreeInstructionTextCallback(BNInstructionTextToken* tokens, size_t count)
+{
+	for (size_t i = 0; i < count; i++)
+		BNFreeString(tokens[i].text);
+	delete[] tokens;
+}
+
+
 void Architecture::Register(Architecture* arch)
 {
 	BNCustomArchitecture callbacks;
 	callbacks.context = arch;
 	callbacks.getInstructionInfo = GetInstructionInfoCallback;
+	callbacks.getInstructionText = GetInstructionTextCallback;
+	callbacks.freeInstructionText = FreeInstructionTextCallback;
 	arch->m_arch = BNRegisterArchitecture(arch->m_nameForRegister.c_str(), &callbacks);
 }
 
@@ -93,3 +136,18 @@ bool CoreArchitecture::GetInstructionInfo(BinaryView* view, uint64_t addr, Instr
 {
 	return BNGetInstructionInfo(m_arch, view->GetViewObject(), addr, &result);
 }
+
+
+bool CoreArchitecture::GetInstructionText(BinaryView* view, uint64_t addr, std::vector<InstructionTextToken>& result)
+{
+	BNInstructionTextToken* tokens = nullptr;
+	size_t count = 0;
+	if (!BNGetInstructionText(m_arch, view->GetViewObject(), addr, &tokens, &count))
+		return false;
+
+	for (size_t i = 0; i < count; i++)
+		result.push_back(InstructionTextToken(tokens[i].type, tokens[i].text, tokens[i].value));
+
+	BNFreeInstructionText(tokens, count);
+	return true;
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 891dcda0..98a481de 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -8,6 +8,7 @@
 #include <vector>
 #include <map>
 #include <exception>
+#include <functional>
 #include "binaryninjacore.h"
 
 #ifdef _MSC_VER
@@ -581,6 +582,16 @@ namespace BinaryNinja
 		void AddBranch(BNBranchType type, uint64_t target = 0);
 	};
 
+	struct InstructionTextToken
+	{
+		BNInstructionTextTokenType type;
+		std::string text;
+		uint64_t value;
+
+		InstructionTextToken();
+		InstructionTextToken(BNInstructionTextTokenType type, const std::string& text, uint64_t value = 0);
+	};
+
 	class Architecture: public RefCountObject
 	{
 	protected:
@@ -590,6 +601,9 @@ namespace BinaryNinja
 		Architecture(BNArchitecture* arch);
 
 		static bool GetInstructionInfoCallback(void* ctxt, BNBinaryView* data, uint64_t addr, BNInstructionInfo* result);
+		static bool GetInstructionTextCallback(void* ctxt, BNBinaryView* data, uint64_t addr,
+		                                       BNInstructionTextToken** result, size_t* count);
+		static void FreeInstructionTextCallback(BNInstructionTextToken* tokens, size_t count);
 
 	public:
 		Architecture(const std::string& name);
@@ -603,6 +617,7 @@ namespace BinaryNinja
 		std::string GetName() const;
 
 		virtual bool GetInstructionInfo(BinaryView* data, uint64_t addr, InstructionInfo& result) = 0;
+		virtual bool GetInstructionText(BinaryView* data, uint64_t addr, std::vector<InstructionTextToken>& result) = 0;
 	};
 
 	class CoreArchitecture: public Architecture
@@ -610,6 +625,7 @@ namespace BinaryNinja
 	public:
 		CoreArchitecture(BNArchitecture* arch);
 		virtual bool GetInstructionInfo(BinaryView* view, uint64_t addr, InstructionInfo& result) override;
+		virtual bool GetInstructionText(BinaryView* view, uint64_t addr, std::vector<InstructionTextToken>& result) override;
 	};
 
 	class Function;
@@ -648,4 +664,61 @@ namespace BinaryNinja
 
 		std::vector<Ref<BasicBlock>> GetBasicBlocks() const;
 	};
+
+	struct FunctionGraphTextLine
+	{
+		uint64_t addr;
+		std::vector<InstructionTextToken> tokens;
+	};
+
+	class FunctionGraphBlock: public RefCountObject
+	{
+		BNFunctionGraphBlock* m_block;
+
+	public:
+		FunctionGraphBlock(BNFunctionGraphBlock* block);
+		~FunctionGraphBlock();
+
+		BNFunctionGraphBlock* GetBlockObject() const { return m_block; }
+
+		int GetX() const;
+		int GetY() const;
+		int GetWidth() const;
+		int GetHeight() const;
+
+		std::vector<FunctionGraphTextLine> GetLines() const;
+		std::vector<BNBasicBlockEdge> GetOutgoingEdges() const;
+	};
+
+	class FunctionGraph: public RefCountObject
+	{
+		BNFunctionGraph* m_graph;
+		std::function<void()> m_completeFunc;
+
+		static void CompleteCallback(void* ctxt);
+
+	public:
+		FunctionGraph(BNFunctionGraph* graph);
+		~FunctionGraph();
+
+		BNFunctionGraph* GetGraphObject() const { return m_graph; }
+
+		Ref<Function> GetFunction() const;
+
+		int GetHorizontalBlockMargin() const;
+		int GetVerticalBlockMargin() const;
+		void SetBlockMargins(int horiz, int vert);
+
+		void StartLayout();
+		bool IsLayoutComplete();
+		void OnComplete(const std::function<void()>& func);
+
+		std::vector<Ref<FunctionGraphBlock>> GetBlocks() const;
+
+		int GetLeftExtent() const;
+		int GetTopExtent() const;
+		int GetRightExtent() const;
+		int GetBottomExtent() const;
+		std::vector<Ref<FunctionGraphBlock>> GetBlocksInRegion(int left, int top, int right, int bottom);
+	};
 }
diff --git a/functiongraph.cpp b/functiongraph.cpp
new file mode 100644
index 00000000..476c2fe2
--- /dev/null
+++ b/functiongraph.cpp
@@ -0,0 +1,117 @@
+#include "binaryninjaapi.h"
+
+using namespace BinaryNinja;
+using namespace std;
+
+
+FunctionGraph::FunctionGraph(BNFunctionGraph* graph): m_graph(graph)
+{
+}
+
+
+FunctionGraph::~FunctionGraph()
+{
+	BNFreeFunctionGraph(m_graph);
+}
+
+
+void FunctionGraph::CompleteCallback(void* ctxt)
+{
+	FunctionGraph* graph = (FunctionGraph*)ctxt;
+	graph->m_completeFunc();
+}
+
+
+Ref<Function> FunctionGraph::GetFunction() const
+{
+	return new Function(BNNewFunctionReference(BNGetFunctionForFunctionGraph(m_graph)));
+}
+
+
+int FunctionGraph::GetHorizontalBlockMargin() const
+{
+	return BNGetHorizontalFunctionGraphBlockMargin(m_graph);
+}
+
+
+int FunctionGraph::GetVerticalBlockMargin() const
+{
+	return BNGetVerticalFunctionGraphBlockMargin(m_graph);
+}
+
+
+void FunctionGraph::SetBlockMargins(int horiz, int vert)
+{
+	BNSetFunctionGraphBlockMargins(m_graph, horiz, vert);
+}
+
+
+void FunctionGraph::StartLayout()
+{
+	BNStartFunctionGraphLayout(m_graph);
+}
+
+
+bool FunctionGraph::IsLayoutComplete()
+{
+	return BNIsFunctionGraphLayoutComplete(m_graph);
+}
+
+
+void FunctionGraph::OnComplete(const std::function<void()>& func)
+{
+	m_completeFunc = func;
+	BNSetFunctionGraphCompleteCallback(m_graph, this, CompleteCallback);
+}
+
+
+vector<Ref<FunctionGraphBlock>> FunctionGraph::GetBlocks() const
+{
+	size_t count;
+	BNFunctionGraphBlock** blocks = BNGetFunctionGraphBlocks(m_graph, &count);
+
+	vector<Ref<FunctionGraphBlock>> result;
+	for (size_t i = 0; i < count; i++)
+		result.push_back(new FunctionGraphBlock(BNNewFunctionGraphBlockReference(blocks[i])));
+
+	BNFreeFunctionGraphBlockList(blocks, count);
+	return result;
+}
+
+
+int FunctionGraph::GetLeftExtent() const
+{
+	return BNGetFunctionGraphLeftExtent(m_graph);
+}
+
+
+int FunctionGraph::GetTopExtent() const
+{
+	return BNGetFunctionGraphTopExtent(m_graph);
+}
+
+
+int FunctionGraph::GetRightExtent() const
+{
+	return BNGetFunctionGraphRightExtent(m_graph);
+}
+
+
+int FunctionGraph::GetBottomExtent() const
+{
+	return BNGetFunctionGraphBottomExtent(m_graph);
+}
+
+
+vector<Ref<FunctionGraphBlock>> FunctionGraph::GetBlocksInRegion(int left, int top, int right, int bottom)
+{
+	size_t count;
+	BNFunctionGraphBlock** blocks = BNGetFunctionGraphBlocksInRegion(m_graph, left, top, right, bottom, &count);
+
+	vector<Ref<FunctionGraphBlock>> result;
+	for (size_t i = 0; i < count; i++)
+		result.push_back(new FunctionGraphBlock(BNNewFunctionGraphBlockReference(blocks[i])));
+
+	BNFreeFunctionGraphBlockList(blocks, count);
+	return result;
+}
diff --git a/functiongraphblock.cpp b/functiongraphblock.cpp
new file mode 100644
index 00000000..8afdef60
--- /dev/null
+++ b/functiongraphblock.cpp
@@ -0,0 +1,78 @@
+#include "binaryninjaapi.h"
+
+using namespace BinaryNinja;
+using namespace std;
+
+
+FunctionGraphBlock::FunctionGraphBlock(BNFunctionGraphBlock* block): m_block(block)
+{
+}
+
+
+FunctionGraphBlock::~FunctionGraphBlock()
+{
+	BNFreeFunctionGraphBlock(m_block);
+}
+
+
+int FunctionGraphBlock::GetX() const
+{
+	return BNGetFunctionGraphBlockX(m_block);
+}
+
+
+int FunctionGraphBlock::GetY() const
+{
+	return BNGetFunctionGraphBlockY(m_block);
+}
+
+
+int FunctionGraphBlock::GetWidth() const
+{
+	return BNGetFunctionGraphBlockWidth(m_block);
+}
+
+
+int FunctionGraphBlock::GetHeight() const
+{
+	return BNGetFunctionGraphBlockHeight(m_block);
+}
+
+
+vector<FunctionGraphTextLine> FunctionGraphBlock::GetLines() const
+{
+	size_t count;
+	BNFunctionGraphTextLine* lines = BNGetFunctionGraphBlockLines(m_block, &count);
+
+	vector<FunctionGraphTextLine> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		FunctionGraphTextLine line;
+		line.addr = lines[i].addr;
+		for (size_t j = 0; j < lines[i].count; j++)
+		{
+			InstructionTextToken token;
+			token.type = lines[i].tokens[j].type;
+			token.text = lines[i].tokens[j].text;
+			token.value = lines[i].tokens[j].value;
+			line.tokens.push_back(token);
+		}
+		result.push_back(line);
+	}
+
+	BNFreeFunctionGraphBlockLines(lines, count);
+	return result;
+}
+
+
+vector<BNBasicBlockEdge> FunctionGraphBlock::GetOutgoingEdges() const
+{
+	size_t count;
+	BNBasicBlockEdge* edges = BNGetFunctionGraphBlockOutgoingEdges(m_block, &count);
+
+	vector<BNBasicBlockEdge> result;
+	result.insert(result.begin(), &edges[0], &edges[count]);
+
+	BNFreeFunctionGraphBlockOutgoingEdgeList(edges);
+	return result;
+}