1 files changed, 31 insertions, 10 deletions

diff --git a/functiongraphblock.cpp b/functiongraphblock.cpp
index 6ad86646..a37c0b49 100644
--- a/functiongraphblock.cpp
+++ b/functiongraphblock.cpp
@@ -27,6 +27,14 @@ using namespace std;
 FunctionGraphBlock::FunctionGraphBlock(BNFunctionGraphBlock* block)
 {
 	m_object = block;
+	m_cachedLinesValid = false;
+	m_cachedEdgesValid = false;
+}
+
+
+Ref<BasicBlock> FunctionGraphBlock::GetBasicBlock() const
+{
+	return new BasicBlock(BNGetFunctionGraphBasicBlock(m_object));
 }
 
 
@@ -72,15 +80,18 @@ int FunctionGraphBlock::GetHeight() const
 }
 
 
-vector<FunctionGraphTextLine> FunctionGraphBlock::GetLines() const
+const vector<DisassemblyTextLine>& FunctionGraphBlock::GetLines()
 {
+	if (m_cachedLinesValid)
+		return m_cachedLines;
+
 	size_t count;
-	BNFunctionGraphTextLine* lines = BNGetFunctionGraphBlockLines(m_object, &count);
+	BNDisassemblyTextLine* lines = BNGetFunctionGraphBlockLines(m_object, &count);
 
-	vector<FunctionGraphTextLine> result;
+	vector<DisassemblyTextLine> result;
 	for (size_t i = 0; i < count; i++)
 	{
-		FunctionGraphTextLine line;
+		DisassemblyTextLine line;
 		line.addr = lines[i].addr;
 		for (size_t j = 0; j < lines[i].count; j++)
 		{
@@ -88,18 +99,27 @@ vector<FunctionGraphTextLine> FunctionGraphBlock::GetLines() const
 			token.type = lines[i].tokens[j].type;
 			token.text = lines[i].tokens[j].text;
 			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);
 	}
 
-	BNFreeFunctionGraphBlockLines(lines, count);
-	return result;
+	BNFreeDisassemblyTextLines(lines, count);
+	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);
 
@@ -108,12 +128,13 @@ vector<FunctionGraphEdge> FunctionGraphBlock::GetOutgoingEdges() const
 	{
 		FunctionGraphEdge edge;
 		edge.type = edges[i].type;
-		edge.target = edges[i].target;
-		edge.arch = edges[i].arch ? new CoreArchitecture(edges[i].arch) : nullptr;
+		edge.target = edges[i].target ? new BasicBlock(BNNewBasicBlockReference(edges[i].target)) : nullptr;
 		edge.points.insert(edge.points.begin(), &edges[i].points[0], &edges[i].points[edges[i].pointCount]);
 		result.push_back(edge);
 	}
 
 	BNFreeFunctionGraphBlockOutgoingEdgeList(edges, count);
-	return result;
+	m_cachedEdges = result;
+	m_cachedEdgesValid = true;
+	return m_cachedEdges;
 }