Share context between BB recovery, lifing, and disassembly

Adds a GetInstructionTextWithContext callback to the architecture class
that can be used to pass data from AnalyzeBasicBlocks. This same context
is also supplied to LiftFunction and allows for supplying shared function
and/or binary view level information across basic block analysis, function
lifting, and disassembly text rendering

1 files changed, 75 insertions, 2 deletions

diff --git a/architecture.cpp b/architecture.cpp
index 1901438f..acc47318 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -307,6 +307,15 @@ std::map<ArchAndAddr, ArchAndAddr>& BasicBlockAnalysisContext::GetInlinedUnresol
 }
 
 
+bool BasicBlockAnalysisContext::SetFunctionArchContextRaw(void* p)
+{
+	if (m_context->functionArchContext)
+		return false;
+	m_context->functionArchContext = p;
+	return true;
+}
+
+
 void BasicBlockAnalysisContext::AddTempOutgoingReference(Function* targetFunc)
 {
 	BNAnalyzeBasicBlocksContextAddTempReference(m_context, targetFunc->m_object);
@@ -461,8 +470,6 @@ void BasicBlockAnalysisContext::Finalize()
 			delete[] values;
 		}
 	}
-
-	BNAnalyzeBasicBlocksContextFinalize(m_context);
 }
 
 
@@ -518,6 +525,7 @@ FunctionLifterContext::FunctionLifterContext(LowLevelILFunction* func, BNFunctio
 		m_inlinedCalls.insert(context->inlinedCalls[i]);
 	}
 
+	m_functionArchContext = context->functionArchContext;
 	m_containsInlinedFunctions = context->containsInlinedFunctions;
 }
 
@@ -730,6 +738,26 @@ bool Architecture::GetInstructionTextCallback(
 }
 
 
+bool Architecture::GetInstructionTextWithContextCallback(void* ctxt, const uint8_t* data, uint64_t addr, size_t* len,
+	void* context, BNInstructionTextToken** result, size_t* count)
+{
+	CallbackRef<Architecture> arch(ctxt);
+
+	vector<InstructionTextToken> tokens;
+	bool ok = arch->GetInstructionTextWithContext(data, addr, *len, context, tokens);
+	if (!ok)
+	{
+		*result = nullptr;
+		*count = 0;
+		return false;
+	}
+
+	*count = tokens.size();
+	*result = InstructionTextToken::CreateInstructionTextTokenList(tokens);
+	return true;
+}
+
+
 void Architecture::FreeInstructionTextCallback(BNInstructionTextToken* tokens, size_t count)
 {
 	for (size_t i = 0; i < count; i++)
@@ -772,6 +800,13 @@ bool Architecture::LiftFunctionCallback(void* ctxt, BNLowLevelILFunction* functi
 }
 
 
+void Architecture::FreeFunctionArchContextCallback(void* ctxt, void* context)
+{
+	CallbackRef<Architecture> arch(ctxt);
+	arch->FreeFunctionArchContext(context);
+}
+
+
 char* Architecture::GetRegisterNameCallback(void* ctxt, uint32_t reg)
 {
 	CallbackRef<Architecture> arch(ctxt);
@@ -1260,10 +1295,12 @@ void Architecture::Register(Architecture* arch)
 	callbacks.getAssociatedArchitectureByAddress = GetAssociatedArchitectureByAddressCallback;
 	callbacks.getInstructionInfo = GetInstructionInfoCallback;
 	callbacks.getInstructionText = GetInstructionTextCallback;
+	callbacks.getInstructionTextWithContext = GetInstructionTextWithContextCallback;
 	callbacks.freeInstructionText = FreeInstructionTextCallback;
 	callbacks.getInstructionLowLevelIL = GetInstructionLowLevelILCallback;
 	callbacks.analyzeBasicBlocks = AnalyzeBasicBlocksCallback;
 	callbacks.liftFunction = LiftFunctionCallback;
+	callbacks.freeFunctionArchContext = FreeFunctionArchContextCallback;
 	callbacks.getRegisterName = GetRegisterNameCallback;
 	callbacks.getFlagName = GetFlagNameCallback;
 	callbacks.getFlagWriteTypeName = GetFlagWriteTypeNameCallback;
@@ -1404,6 +1441,16 @@ bool Architecture::LiftFunction(LowLevelILFunction* function, FunctionLifterCont
 }
 
 
+bool Architecture::GetInstructionTextWithContext(
+	const uint8_t* data, uint64_t addr, size_t& len, void* context, std::vector<InstructionTextToken>& result)
+{
+	return GetInstructionText(data, addr, len, result);
+}
+
+
+void Architecture::FreeFunctionArchContext(void* context) {}
+
+
 string Architecture::GetRegisterName(uint32_t reg)
 {
 	return fmt::format("r{}", reg);
@@ -1959,6 +2006,19 @@ bool CoreArchitecture::GetInstructionText(
 }
 
 
+bool CoreArchitecture::GetInstructionTextWithContext(
+	const uint8_t* data, uint64_t addr, size_t& len, void* context, std::vector<InstructionTextToken>& result)
+{
+	BNInstructionTextToken* tokens = nullptr;
+	size_t count = 0;
+	if (!BNGetInstructionTextWithContext(m_object, data, addr, &len, context, &tokens, &count))
+		return false;
+
+	result = InstructionTextToken::ConvertAndFreeInstructionTextTokenList(tokens, count);
+	return true;
+}
+
+
 bool CoreArchitecture::GetInstructionLowLevelIL(const uint8_t* data, uint64_t addr, size_t& len, LowLevelILFunction& il)
 {
 	return BNGetInstructionLowLevelIL(m_object, data, addr, &len, il.GetObject());
@@ -1977,6 +2037,12 @@ bool CoreArchitecture::LiftFunction(LowLevelILFunction* function, FunctionLifter
 }
 
 
+void CoreArchitecture::FreeFunctionArchContext(void* context)
+{
+	BNArchitectureFreeFunctionArchContext(m_object, context);
+}
+
+
 string CoreArchitecture::GetRegisterName(uint32_t reg)
 {
 	char* name = BNGetArchitectureRegisterName(m_object, reg);
@@ -2487,6 +2553,13 @@ bool ArchitectureExtension::GetInstructionText(
 }
 
 
+bool ArchitectureExtension::GetInstructionTextWithContext(
+	const uint8_t* data, uint64_t addr, size_t& len, void* context, vector<InstructionTextToken>& result)
+{
+	return m_base->GetInstructionTextWithContext(data, addr, len, context, result);
+}
+
+
 bool ArchitectureExtension::GetInstructionLowLevelIL(
     const uint8_t* data, uint64_t addr, size_t& len, LowLevelILFunction& il)
 {