Support MLIL expr mappings in C++

6 files changed, 273 insertions, 33 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 344fb03e..32d6051f 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -11106,11 +11106,20 @@ namespace BinaryNinja {
 		*/
 		void SetLowLevelILFunction(Ref<LowLevelILFunction> lowLevelIL);
 
-		/*! Set the new Medium Level IL for the current analysis context
+		/*! Set the new Medium Level IL for the current analysis context.
+
+			If mapping parameters are left as default (empty), then they will be automatically
+			computed for you based on previous calls to AddExpr() and AddInstruction()
 
 			\param mediumLevelIL the new Medium Level IL
+			\param llilSsaToMlilInstrMap New mappings from LLIL SSA -> MLIL instruction indices
+			\param llilSsaToMlilExprMap New mappings from LLIL SSA -> MLIL expression indices
 		*/
-		void SetMediumLevelILFunction(Ref<MediumLevelILFunction> mediumLevelIL);
+		void SetMediumLevelILFunction(
+			Ref<MediumLevelILFunction> mediumLevelIL,
+			std::unordered_map<size_t /* llil ssa */, size_t /* mlil */> llilSsaToMlilInstrMap = {},
+			std::vector<BNExprMapInfo> llilSsaToMlilExprMap = {}
+		);
 
 		/*! Set the new High Level IL for the current analysis context
 
@@ -13071,21 +13080,16 @@ namespace BinaryNinja {
 		uint32_t sourceOperand;
 		bool valid;
 
-		ILSourceLocation() : valid(false) {}
+		bool ilBased;
+		bool ilDirect;
+		size_t ilExprIndex;
 
-		ILSourceLocation(uint64_t addr, uint32_t operand) : address(addr), sourceOperand(operand), valid(true) {}
-
-		ILSourceLocation(const BNLowLevelILInstruction& instr) :
-		    address(instr.address), sourceOperand(instr.sourceOperand), valid(true)
-		{}
+		ILSourceLocation() : valid(false), ilBased(false) {}
 
-		ILSourceLocation(const BNMediumLevelILInstruction& instr) :
-		    address(instr.address), sourceOperand(instr.sourceOperand), valid(true)
-		{}
-
-		ILSourceLocation(const BNHighLevelILInstruction& instr) :
-		    address(instr.address), sourceOperand(instr.sourceOperand), valid(true)
-		{}
+		ILSourceLocation(uint64_t addr, uint32_t operand) : address(addr), sourceOperand(operand), valid(true), ilBased(false) {}
+		ILSourceLocation(const struct LowLevelILInstruction& instr);
+		ILSourceLocation(const struct MediumLevelILInstruction& instr);
+		ILSourceLocation(const struct HighLevelILInstruction& instr);
 	};
 
 	struct LowLevelILInstruction;
@@ -14535,6 +14539,22 @@ namespace BinaryNinja {
 	    public CoreRefCountObject<BNMediumLevelILFunction, BNNewMediumLevelILFunctionReference,
 	        BNFreeMediumLevelILFunction>
 	{
+		struct TranslationData
+		{
+			MediumLevelILFunction* copyingFunction = nullptr;
+			std::unordered_map<size_t /* old function expr index */, std::vector<std::tuple<size_t /* new function expr index */, bool /* direct */>>> mlilToMlilExprMap;
+			std::unordered_map<size_t /* old function instr index */, std::vector<std::tuple<size_t /* new function instr index */, bool /* direct */>>> mlilToMlilInstrMap;
+			// todo maybe: llil ssa -> mlil mappings
+		};
+		std::unique_ptr<TranslationData> m_translationData;
+
+		void RecordMLILToMLILExprMap(size_t newExprIndex, const ILSourceLocation& location);
+		void RecordMLILToMLILInstrMap(size_t newInstrIndex, const ILSourceLocation& location);
+		std::unordered_map<size_t /* llil ssa */, size_t /* mlil */> GetLLILSSAToMLILInstrMap(bool fromTranslation);
+		std::vector<BNExprMapInfo> GetLLILSSAToMLILExprMap(bool fromTranslation);
+
+		friend class AnalysisContext;
+
 	  public:
 		MediumLevelILFunction(Architecture* arch, Function* func = nullptr, LowLevelILFunction* lowLevelIL = nullptr);
 		MediumLevelILFunction(BNMediumLevelILFunction* func);
@@ -14774,7 +14794,7 @@ namespace BinaryNinja {
 		    const ILSourceLocation& loc = ILSourceLocation());
 		void MarkLabel(BNMediumLevelILLabel& label);
 
-		ExprId AddInstruction(ExprId expr);
+		ExprId AddInstruction(ExprId expr, const ILSourceLocation& loc = ILSourceLocation());
 
 		std::vector<uint64_t> GetOperandList(ExprId i, size_t listOperand);
 		ExprId AddLabelMap(const std::map<uint64_t, BNMediumLevelILLabel*>& labels);
diff --git a/examples/workflows/unflatten/library.cpp b/examples/workflows/unflatten/library.cpp
index 062b928e..cb4e9b9a 100644
--- a/examples/workflows/unflatten/library.cpp
+++ b/examples/workflows/unflatten/library.cpp
@@ -235,7 +235,7 @@ void RewriteAction(Ref<AnalysisContext> context, bool doIt)
 						{
 							node->SetHighlight(GetHighlightColor(RedHighlightColor));
 						}
-						else if (std::find_if(path.begin(), path.end(), [node](Ref<BasicBlock> b) { return b->GetStart() == node->GetBasicBlock()->GetStart(); }) != path.end())
+						else if (std::find_if(path.begin(), path.end(), [node](const Ref<BasicBlock>& b) { return b->GetStart() == node->GetBasicBlock()->GetStart(); }) != path.end())
 						{
 							node->SetHighlight(GetHighlightColor(GreenHighlightColor));
 						}
@@ -286,7 +286,7 @@ void RewriteAction(Ref<AnalysisContext> context, bool doIt)
 									// Copy instruction as-is
 									auto copyBlockInstr = oldMLIL->GetInstruction(copyBlockInstrIndex);
 									newMLIL->SetCurrentAddress(copyBlock->GetArchitecture(), copyBlockInstr.address);
-									newMLIL->AddInstruction(copyBlockInstr.CopyTo(newMLIL));
+									newMLIL->AddInstruction(copyBlockInstr.CopyTo(newMLIL), copyBlockInstr);
 								}
 							}
 							continue;
@@ -295,7 +295,7 @@ void RewriteAction(Ref<AnalysisContext> context, bool doIt)
 
 					// Otherwise, copy the instruction as-is
 					newMLIL->SetCurrentAddress(block->GetArchitecture(), oldInstr.address);
-					newMLIL->AddInstruction(oldInstr.CopyTo(newMLIL));
+					newMLIL->AddInstruction(oldInstr.CopyTo(newMLIL), oldInstr);
 				}
 			}
 
@@ -350,7 +350,7 @@ void RewriteAction(Ref<AnalysisContext> context, bool doIt)
 						{
 							size_t destValue = oldInstr.GetDestExpr<MLIL_JUMP_TO>().GetValue().value;
 							auto targets = oldInstr.GetTargets<MLIL_JUMP_TO>();
-							if (std::find_if(targets.begin(), targets.end(), [&](std::pair<size_t, size_t> target) {
+							if (std::find_if(targets.begin(), targets.end(), [&](const std::pair<size_t, size_t>& target) {
 								return target.first == destValue;
 							}) != targets.end()) {
 								auto oldTargetIndex = targets[destValue];
@@ -359,7 +359,7 @@ void RewriteAction(Ref<AnalysisContext> context, bool doIt)
 									blockLabels[oldTargetIndex] = MediumLevelILLabel{};
 								}
 								MediumLevelILLabel* targetLabel = &blockLabels[oldTargetIndex];
-								newMLIL->AddInstruction(newMLIL->Goto(*targetLabel, oldInstr));
+								newMLIL->AddInstruction(newMLIL->Goto(*targetLabel, oldInstr), oldInstr);
 								continue;
 							}
 						}
@@ -367,7 +367,7 @@ void RewriteAction(Ref<AnalysisContext> context, bool doIt)
 
 					// Otherwise, copy the instruction as-is
 					newMLIL->SetCurrentAddress(block->GetArchitecture(), oldInstr.address);
-					newMLIL->AddInstruction(oldInstr.CopyTo(newMLIL));
+					newMLIL->AddInstruction(oldInstr.CopyTo(newMLIL), oldInstr);
 				}
 			}
 
diff --git a/highlevelil.cpp b/highlevelil.cpp
index 77a0f051..84d74eb3 100644
--- a/highlevelil.cpp
+++ b/highlevelil.cpp
@@ -26,6 +26,12 @@ using namespace BinaryNinja;
 using namespace std;
 
 
+ILSourceLocation::ILSourceLocation(const struct HighLevelILInstruction& instr):
+	address(instr.address), sourceOperand(instr.sourceOperand), valid(true),
+	ilBased(true), ilDirect(true), ilExprIndex(instr.exprIndex)
+{}
+
+
 HighLevelILFunction::HighLevelILFunction(Architecture* arch, Function* func)
 {
 	m_object = BNCreateHighLevelILFunction(arch->GetObject(), func ? func->GetObject() : nullptr);
diff --git a/lowlevelil.cpp b/lowlevelil.cpp
index dd3989f1..5fe2ed91 100644
--- a/lowlevelil.cpp
+++ b/lowlevelil.cpp
@@ -25,6 +25,12 @@ using namespace BinaryNinja;
 using namespace std;
 
 
+ILSourceLocation::ILSourceLocation(const struct LowLevelILInstruction& instr):
+	address(instr.address), sourceOperand(instr.sourceOperand), valid(true),
+	ilBased(true), ilDirect(true), ilExprIndex(instr.exprIndex)
+{}
+
+
 LowLevelILLabel::LowLevelILLabel()
 {
 	BNLowLevelILInitLabel(this);
diff --git a/mediumlevelil.cpp b/mediumlevelil.cpp
index 4f6afe27..26a020ae 100644
--- a/mediumlevelil.cpp
+++ b/mediumlevelil.cpp
@@ -25,12 +25,162 @@ using namespace BinaryNinja;
 using namespace std;
 
 
+ILSourceLocation::ILSourceLocation(const struct MediumLevelILInstruction& instr):
+	address(instr.address), sourceOperand(instr.sourceOperand), valid(true),
+	ilBased(true), ilDirect(true), ilExprIndex(instr.exprIndex)
+{}
+
+
 MediumLevelILLabel::MediumLevelILLabel()
 {
 	BNMediumLevelILInitLabel(this);
 }
 
 
+void MediumLevelILFunction::RecordMLILToMLILExprMap(size_t newExprIndex, const ILSourceLocation& location)
+{
+	if (m_translationData && m_translationData->copyingFunction && location.valid && location.ilBased)
+	{
+		size_t oldExprIndex = location.ilExprIndex;
+		if (m_translationData->mlilToMlilExprMap.find(oldExprIndex) == m_translationData->mlilToMlilExprMap.end())
+		{
+			m_translationData->mlilToMlilExprMap.insert({oldExprIndex, {}});
+		}
+		m_translationData->mlilToMlilExprMap[oldExprIndex].push_back({ newExprIndex, location.ilDirect });
+	}
+}
+
+
+void MediumLevelILFunction::RecordMLILToMLILInstrMap(size_t newInstrIndex, const ILSourceLocation& location)
+{
+	if (m_translationData && m_translationData->copyingFunction && location.valid && location.ilBased)
+	{
+		size_t oldExprIndex = location.ilExprIndex;
+		size_t oldInstrIndex = m_translationData->copyingFunction->GetInstructionForExpr(oldExprIndex);
+		if (m_translationData->mlilToMlilInstrMap.find(oldInstrIndex) == m_translationData->mlilToMlilInstrMap.end())
+		{
+			m_translationData->mlilToMlilInstrMap.insert({oldInstrIndex, {}});
+		}
+		m_translationData->mlilToMlilInstrMap[oldInstrIndex].push_back({ newInstrIndex, location.ilDirect });
+	}
+}
+
+
+std::unordered_map<size_t /* llil ssa */, size_t /* mlil */> MediumLevelILFunction::GetLLILSSAToMLILInstrMap(bool fromTranslation)
+{
+	std::unordered_map<size_t /* llil ssa */, size_t /* mlil */> result;
+	if (fromTranslation)
+	{
+		// TODO: Handle LLIL SSA -> MLIL mappings in case someone is brave enough to try
+		// lifting LLILSSA->MLIL themselves instead of an MLIL->MLIL translation
+		// (which is the only one I've seen people do so far)
+
+		if (m_translationData && m_translationData->copyingFunction)
+		{
+			for (auto& [oldInstrIndex, newInstrIndices]: m_translationData->mlilToMlilInstrMap)
+			{
+				// Look up the LLIL SSA instruction for the old instr in its function
+				// and then store that mapping for the new function
+
+				for (auto& [newInstrIndex, newDirect]: newInstrIndices)
+				{
+					// Instructions are always mapped 1 to 1. If the map is marked indirect
+					// then just ignore it.
+					if (newDirect)
+					{
+						size_t oldLLILSSAIndex = m_translationData->copyingFunction->GetLowLevelILInstructionIndex(oldInstrIndex);
+						if (oldLLILSSAIndex != BN_INVALID_EXPR)
+						{
+							result[oldLLILSSAIndex] = newInstrIndex;
+						}
+					}
+				}
+			}
+		}
+	}
+	else
+	{
+		for (auto& block: GetBasicBlocks())
+		{
+			for (size_t instrIndex = block->GetStart(); instrIndex < block->GetEnd(); instrIndex++)
+			{
+				size_t llilSSAIndex = GetLowLevelILInstructionIndex(instrIndex);
+				result[llilSSAIndex] = instrIndex;
+			}
+		}
+	}
+
+	return result;
+}
+
+
+std::vector<BNExprMapInfo> MediumLevelILFunction::GetLLILSSAToMLILExprMap(bool fromTranslation)
+{
+	std::vector<BNExprMapInfo> result;
+	if (fromTranslation)
+	{
+		// TODO: Handle LLIL SSA -> MLIL mappings in case someone is brave enough to try
+		// lifting LLILSSA->MLIL themselves instead of an MLIL->MLIL translation
+		// (which is the only one I've seen people do so far)
+
+		for (auto& [oldExprIndex, newExprIndices]: m_translationData->mlilToMlilExprMap)
+		{
+			// Look up the LLIL SSA expression for the old expr in its function
+			// And then store that mapping for the new function
+
+			size_t oldLLILSSADirect = m_translationData->copyingFunction->GetLowLevelILExprIndex(oldExprIndex);
+			auto oldLLILSSAIndices = m_translationData->copyingFunction->GetLowLevelILExprIndexes(oldExprIndex);
+			for (auto& oldLLILSSAIndex: oldLLILSSAIndices)
+			{
+				size_t oldReverseDirect = m_translationData->copyingFunction->GetLowLevelIL()->GetSSAForm()->GetMediumLevelILExprIndex(oldLLILSSAIndex);
+				auto oldReverseAll = m_translationData->copyingFunction->GetLowLevelIL()->GetSSAForm()->GetMediumLevelILExprIndexes(oldLLILSSAIndex);
+				for (auto& [newExprIndex, newDirect]: newExprIndices)
+				{
+					BNExprMapInfo info;
+					info.lowerIndex = oldLLILSSAIndex;
+					info.higherIndex = newExprIndex;
+					info.lowerToHigherDirect = newDirect && oldReverseDirect == oldExprIndex;
+					info.higherToLowerDirect = newDirect && oldExprIndex == oldLLILSSADirect;
+					info.mapLowerToHigher = oldReverseAll.contains(oldExprIndex);
+					info.mapHigherToLower = true;
+					result.push_back(info);
+				}
+			}
+		}
+	}
+	else
+	{
+		for (auto& block: GetBasicBlocks())
+		{
+			for (size_t instrIndex = block->GetStart(); instrIndex < block->GetEnd(); instrIndex++)
+			{
+				GetInstruction(instrIndex).VisitExprs([&](const MediumLevelILInstruction& expr)
+				{
+					size_t llilSSADirect = GetLowLevelILExprIndex(expr.exprIndex);
+					auto llilSSAIndices = GetLowLevelILExprIndexes(expr.exprIndex);
+					for (auto& llilSSAIndex: llilSSAIndices)
+					{
+						size_t reverseDirect = GetLowLevelIL()->GetSSAForm()->GetMediumLevelILExprIndex(llilSSAIndex);
+						auto reverseAll = GetLowLevelIL()->GetSSAForm()->GetMediumLevelILExprIndexes(llilSSAIndex);
+
+						BNExprMapInfo info;
+						info.lowerIndex = llilSSAIndex;
+						info.higherIndex = expr.exprIndex;
+						info.lowerToHigherDirect = reverseDirect == expr.exprIndex;
+						info.higherToLowerDirect = llilSSAIndex == llilSSADirect;
+						info.mapLowerToHigher = reverseAll.contains(expr.exprIndex);
+						info.mapHigherToLower = true;
+						result.push_back(info);
+					}
+					return true;
+				});
+			}
+		}
+	}
+	return result;
+}
+
+
 MediumLevelILFunction::MediumLevelILFunction(Architecture* arch, Function* func, LowLevelILFunction* lowLevelIL)
 {
 	m_object = BNCreateMediumLevelILFunction(arch->GetObject(), func ? func->GetObject() : nullptr, lowLevelIL ? lowLevelIL->GetObject() : nullptr);
@@ -81,6 +231,11 @@ size_t MediumLevelILFunction::GetInstructionStart(Architecture* arch, uint64_t a
 
 void MediumLevelILFunction::PrepareToCopyFunction(MediumLevelILFunction* func)
 {
+	if (!m_translationData)
+	{
+		m_translationData = std::make_unique<MediumLevelILFunction::TranslationData>();
+	}
+	m_translationData->copyingFunction = func;
 	BNPrepareToCopyMediumLevelILFunction(m_object, func->GetObject());
 }
 
@@ -128,35 +283,59 @@ ExprId MediumLevelILFunction::AddExprWithLocation(BNMediumLevelILOperation opera
 ExprId MediumLevelILFunction::AddExprWithLocation(BNMediumLevelILOperation operation, const ILSourceLocation& loc,
     size_t size, ExprId a, ExprId b, ExprId c, ExprId d, ExprId e)
 {
+	ExprId index;
 	if (loc.valid)
 	{
-		return BNMediumLevelILAddExprWithLocation(
+		index = BNMediumLevelILAddExprWithLocation(
 		    m_object, operation, loc.address, loc.sourceOperand, size, a, b, c, d, e);
 	}
-	return BNMediumLevelILAddExpr(m_object, operation, size, a, b, c, d, e);
+	else
+	{
+		index = BNMediumLevelILAddExpr(m_object, operation, size, a, b, c, d, e);
+	}
+	RecordMLILToMLILExprMap(index, loc);
+	return index;
 }
 
 
-ExprId MediumLevelILFunction::AddInstruction(size_t expr)
+ExprId MediumLevelILFunction::AddInstruction(size_t expr, const ILSourceLocation& loc)
 {
-	return BNMediumLevelILAddInstruction(m_object, expr);
+	ExprId index = BNMediumLevelILAddInstruction(m_object, expr);
+	RecordMLILToMLILInstrMap(index, loc);
+	return index;
 }
 
 
 ExprId MediumLevelILFunction::Goto(BNMediumLevelILLabel& label, const ILSourceLocation& loc)
 {
+	size_t index;
 	if (loc.valid)
-		return BNMediumLevelILGotoWithLocation(m_object, &label, loc.address, loc.sourceOperand);
-	return BNMediumLevelILGoto(m_object, &label);
+	{
+		index = BNMediumLevelILGotoWithLocation(m_object, &label, loc.address, loc.sourceOperand);
+	}
+	else
+	{
+		index = BNMediumLevelILGoto(m_object, &label);
+	}
+	RecordMLILToMLILExprMap(index, loc);
+	return index;
 }
 
 
 ExprId MediumLevelILFunction::If(
     ExprId operand, BNMediumLevelILLabel& t, BNMediumLevelILLabel& f, const ILSourceLocation& loc)
 {
+	size_t index;
 	if (loc.valid)
-		return BNMediumLevelILIfWithLocation(m_object, operand, &t, &f, loc.address, loc.sourceOperand);
-	return BNMediumLevelILIf(m_object, operand, &t, &f);
+	{
+		index = BNMediumLevelILIfWithLocation(m_object, operand, &t, &f, loc.address, loc.sourceOperand);
+	}
+	else
+	{
+		index = BNMediumLevelILIf(m_object, operand, &t, &f);
+	}
+	RecordMLILToMLILExprMap(index, loc);
+	return index;
 }
 
 
diff --git a/workflow.cpp b/workflow.cpp
index 03a10c54..13c180a5 100644
--- a/workflow.cpp
+++ b/workflow.cpp
@@ -98,10 +98,39 @@ void AnalysisContext::SetLowLevelILFunction(Ref<LowLevelILFunction> lowLevelIL)
 }
 
 
-void AnalysisContext::SetMediumLevelILFunction(Ref<MediumLevelILFunction> mediumLevelIL)
+void AnalysisContext::SetMediumLevelILFunction(
+	Ref<MediumLevelILFunction> mediumLevelIL,
+	std::unordered_map<size_t /* llil ssa */, size_t /* mlil */> llilSsaToMlilInstrMap,
+	std::vector<BNExprMapInfo> llilSsaToMlilExprMap
+)
 {
-	// TODO: Mappings FFI
-	BNSetMediumLevelILFunction(m_object, mediumLevelIL->m_object, nullptr, 0, nullptr, 0);
+	if (llilSsaToMlilExprMap.empty() || llilSsaToMlilInstrMap.empty())
+	{
+		// Build up maps from existing data in the function
+		llilSsaToMlilExprMap = mediumLevelIL->GetLLILSSAToMLILExprMap(true);
+		llilSsaToMlilInstrMap = mediumLevelIL->GetLLILSSAToMLILInstrMap(true);
+	}
+
+	std::vector<size_t> instrMapVec;
+	// Technically imprecise but doesn't matter this is just reserving
+	instrMapVec.reserve(llilSsaToMlilInstrMap.size());
+	for (auto& [llilSSAIndex, mlilIndex]: llilSsaToMlilInstrMap)
+	{
+		if (instrMapVec.size() <= llilSSAIndex)
+		{
+			instrMapVec.resize(llilSSAIndex + 1, BN_INVALID_EXPR);
+		}
+		instrMapVec[llilSSAIndex] = mlilIndex;
+	}
+
+	BNSetMediumLevelILFunction(
+		m_object,
+		mediumLevelIL->m_object,
+		instrMapVec.data(),
+		instrMapVec.size(),
+		llilSsaToMlilExprMap.data(),
+		llilSsaToMlilExprMap.size()
+	);
 }