Cache lines on API side to avoid UI hangs with large basic blocks

3 files changed, 51 insertions, 10 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 6902c1a3..fed1bbd9 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1694,6 +1694,10 @@ namespace BinaryNinja
 	class FunctionGraphBlock: public CoreRefCountObject<BNFunctionGraphBlock,
 		BNNewFunctionGraphBlockReference, BNFreeFunctionGraphBlock>
 	{
+		std::vector<DisassemblyTextLine> m_cachedLines;
+		std::vector<FunctionGraphEdge> m_cachedEdges;
+		bool m_cachedLinesValid, m_cachedEdgesValid;
+
 	public:
 		FunctionGraphBlock(BNFunctionGraphBlock* block);
 
@@ -1705,14 +1709,15 @@ namespace BinaryNinja
 		int GetWidth() const;
 		int GetHeight() const;
 
-		std::vector<DisassemblyTextLine> GetLines() const;
-		std::vector<FunctionGraphEdge> GetOutgoingEdges() const;
+		const std::vector<DisassemblyTextLine>& GetLines();
+		const std::vector<FunctionGraphEdge>& GetOutgoingEdges();
 	};
 
 	class FunctionGraph: public RefCountObject
 	{
 		BNFunctionGraph* m_graph;
 		std::function<void()> m_completeFunc;
+		std::map<BNFunctionGraphBlock*, Ref<FunctionGraphBlock>> m_cachedBlocks;
 
 		static void CompleteCallback(void* ctxt);
 
@@ -1735,7 +1740,7 @@ namespace BinaryNinja
 		void OnComplete(const std::function<void()>& func);
 		void Abort();
 
-		std::vector<Ref<FunctionGraphBlock>> GetBlocks() const;
+		std::vector<Ref<FunctionGraphBlock>> GetBlocks();
 
 		int GetWidth() const;
 		int GetHeight() const;
diff --git a/functiongraph.cpp b/functiongraph.cpp
index 93f60f63..30e2e165 100644
--- a/functiongraph.cpp
+++ b/functiongraph.cpp
@@ -104,14 +104,26 @@ void FunctionGraph::Abort()
 }
 
 
-vector<Ref<FunctionGraphBlock>> FunctionGraph::GetBlocks() const
+vector<Ref<FunctionGraphBlock>> FunctionGraph::GetBlocks()
 {
 	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])));
+	{
+		auto block = m_cachedBlocks.find(blocks[i]);
+		if (block == m_cachedBlocks.end())
+		{
+			FunctionGraphBlock* newBlock = new FunctionGraphBlock(BNNewFunctionGraphBlockReference(blocks[i]));
+			m_cachedBlocks[blocks[i]] = newBlock;
+			result.push_back(newBlock);
+		}
+		else
+		{
+			result.push_back(block->second);
+		}
+	}
 
 	BNFreeFunctionGraphBlockList(blocks, count);
 	return result;
@@ -137,7 +149,19 @@ vector<Ref<FunctionGraphBlock>> FunctionGraph::GetBlocksInRegion(int left, int t
 
 	vector<Ref<FunctionGraphBlock>> result;
 	for (size_t i = 0; i < count; i++)
-		result.push_back(new FunctionGraphBlock(BNNewFunctionGraphBlockReference(blocks[i])));
+	{
+		auto block = m_cachedBlocks.find(blocks[i]);
+		if (block == m_cachedBlocks.end())
+		{
+			FunctionGraphBlock* newBlock = new FunctionGraphBlock(BNNewFunctionGraphBlockReference(blocks[i]));
+			m_cachedBlocks[blocks[i]] = newBlock;
+			result.push_back(newBlock);
+		}
+		else
+		{
+			result.push_back(block->second);
+		}
+	}
 
 	BNFreeFunctionGraphBlockList(blocks, count);
 	return result;
diff --git a/functiongraphblock.cpp b/functiongraphblock.cpp
index af330497..436339f6 100644
--- a/functiongraphblock.cpp
+++ b/functiongraphblock.cpp
@@ -27,6 +27,8 @@ using namespace std;
 FunctionGraphBlock::FunctionGraphBlock(BNFunctionGraphBlock* block)
 {
 	m_object = block;
+	m_cachedLinesValid = false;
+	m_cachedEdgesValid = false;
 }
 
 
@@ -72,8 +74,11 @@ int FunctionGraphBlock::GetHeight() const
 }
 
 
-vector<DisassemblyTextLine> FunctionGraphBlock::GetLines() const
+const vector<DisassemblyTextLine>& FunctionGraphBlock::GetLines()
 {
+	if (m_cachedLinesValid)
+		return m_cachedLines;
+
 	size_t count;
 	BNDisassemblyTextLine* lines = BNGetFunctionGraphBlockLines(m_object, &count);
 
@@ -96,12 +101,17 @@ vector<DisassemblyTextLine> FunctionGraphBlock::GetLines() const
 	}
 
 	BNFreeDisassemblyTextLines(lines, count);
-	return result;
+	m_cachedLines = result;
+	m_cachedLinesValid = true;
+	return m_cachedLines;
 }
 
 
-vector<FunctionGraphEdge> FunctionGraphBlock::GetOutgoingEdges() const
+const vector<FunctionGraphEdge>& FunctionGraphBlock::GetOutgoingEdges()
 {
+	if (m_cachedEdgesValid)
+		return m_cachedEdges;
+
 	size_t count;
 	BNFunctionGraphEdge* edges = BNGetFunctionGraphBlockOutgoingEdges(m_object, &count);
 
@@ -117,5 +127,7 @@ vector<FunctionGraphEdge> FunctionGraphBlock::GetOutgoingEdges() const
 	}
 
 	BNFreeFunctionGraphBlockOutgoingEdgeList(edges, count);
-	return result;
+	m_cachedEdges = result;
+	m_cachedEdgesValid = true;
+	return m_cachedEdges;
 }