Memoize reverse expr lookups when building the LLIL-SSA <-> MLIL expr map

`GetMediumLevelILExprIndexes` now derives its result by walking the
LLIL-SSA use-def graph rather than reading precomputed data.
`GetLLILSSAToMLILExprMap` queries it once per (MLIL expr, LLIL-SSA expr)
pair, and the same LLIL-SSA expr recurs across many MLIL exprs, so a
whole-function map build repeats the same walk many times over. Cache
each expr's reverse result for the life of one map build.

Author: bdash

mediumlevelil.cpp

@@ -24,6 +24,37 @@
 using namespace BinaryNinja;
 using namespace std;
 
+namespace {
+
+// Memoizes the reverse (LLIL-SSA -> MLIL) expr lookups for one LLIL-SSA function.
+// The same LLIL-SSA expr is reached from many MLIL exprs, and a reverse query
+// derives its result by walking the LLIL-SSA use-def graph rather than reading a
+// stored table, so repeating it for each recurrence is what dominates a
+// whole-function map build.
+class ReverseExprMemo
+{
+	Ref<LowLevelILFunction> m_llilSsa;
+	std::unordered_map<size_t, std::pair<size_t, std::set<size_t>>> m_entries;
+
+public:
+	explicit ReverseExprMemo(Ref<LowLevelILFunction> llilSsa) : m_llilSsa(std::move(llilSsa)) {}
+
+	// The MLIL expr that `llilSsaExpr` maps to directly, paired with every MLIL expr it maps to.
+	const std::pair<size_t, std::set<size_t>>& Get(size_t llilSsaExpr)
+	{
+		auto entry = m_entries.find(llilSsaExpr);
+		if (entry == m_entries.end())
+		{
+			entry = m_entries.emplace(llilSsaExpr,
+				std::make_pair(m_llilSsa->GetMediumLevelILExprIndex(llilSsaExpr), m_llilSsa->GetMediumLevelILExprIndexes(llilSsaExpr))
+			).first;
+		}
+		return entry->second;
+	}
+};
+
+}  // unnamed namespace
+
 
 ILSourceLocation::ILSourceLocation(const struct MediumLevelILInstruction& instr):
 	address(instr.address), sourceOperand(instr.sourceOperand), valid(true),
@@ -117,12 +148,14 @@ std::unordered_map<size_t /* llil ssa */, size_t /* mlil */> MediumLevelILFuncti
 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)
 
+		ReverseExprMemo reverse(m_translationData->copyingFunction->GetLowLevelIL()->GetSSAForm());
 		for (auto& [oldExprIndex, newExprIndices]: m_translationData->mlilToMlilExprMap)
 		{
 			// Look up the LLIL SSA expression for the old expr in its function
@@ -132,8 +165,7 @@ std::vector<BNExprMapInfo> MediumLevelILFunction::GetLLILSSAToMLILExprMap(bool f
 			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);
+				const auto& [oldReverseDirect, oldReverseAll] = reverse.Get(oldLLILSSAIndex);
 				for (auto& [newExprIndex, newDirect]: newExprIndices)
 				{
 					BNExprMapInfo info;
@@ -150,6 +182,7 @@ std::vector<BNExprMapInfo> MediumLevelILFunction::GetLLILSSAToMLILExprMap(bool f
 	}
 	else
 	{
+		ReverseExprMemo reverse(GetLowLevelIL()->GetSSAForm());
 		for (auto& block: GetBasicBlocks())
 		{
 			for (size_t instrIndex = block->GetStart(); instrIndex < block->GetEnd(); instrIndex++)
@@ -160,8 +193,7 @@ std::vector<BNExprMapInfo> MediumLevelILFunction::GetLLILSSAToMLILExprMap(bool f
 					auto llilSSAIndices = GetLowLevelILExprIndexes(expr.exprIndex);
 					for (auto& llilSSAIndex: llilSSAIndices)
 					{
-						size_t reverseDirect = GetLowLevelIL()->GetSSAForm()->GetMediumLevelILExprIndex(llilSSAIndex);
-						auto reverseAll = GetLowLevelIL()->GetSSAForm()->GetMediumLevelILExprIndexes(llilSSAIndex);
+						const auto& [reverseDirect, reverseAll] = reverse.Get(llilSSAIndex);
 
 						BNExprMapInfo info;
 						info.lowerIndex = llilSSAIndex;

python/mediumlevelil.py

@@ -4511,6 +4511,24 @@ def _get_llil_ssa_to_mlil_instr_map(self, from_builders: bool) -> LLILSSAToMLILI
 	def _get_llil_ssa_to_mlil_expr_map(self, from_builders: bool) -> LLILSSAToMLILExpressionMapping:
 		llil_ssa_to_mlil_expr_map = []
 
+		# Memoize the reverse (LLIL-SSA -> MLIL) expr lookups, per LLIL-SSA function. The same LLIL-SSA expr is
+		# reached from many MLIL exprs, and a reverse query derives its result by walking the LLIL-SSA use-def graph
+		# rather than reading a stored table, so repeating it for each recurrence is what dominates a whole-function
+		# map build.
+		reverse_memo: Dict['lowlevelil.LowLevelILFunction', Dict[ExpressionIndex, Tuple[
+			Optional[ExpressionIndex], List[ExpressionIndex]]]] = {}
+
+		def reverse_lookup(llil_ssa, llil_ssa_expr):
+			entries = reverse_memo.setdefault(llil_ssa, {})
+			entry = entries.get(llil_ssa_expr)
+			if entry is None:
+				entry = (
+					llil_ssa.get_medium_level_il_expr_index(llil_ssa_expr),
+					llil_ssa.get_medium_level_il_expr_indexes(llil_ssa_expr)
+				)
+				entries[llil_ssa_expr] = entry
+			return entry
+
 		if from_builders:
 			# TODO: Handle LLIL SSA -> MLIL mappings in case someone is brave enough to try
 			# lifting LLILSSA->MLIL themselves instead of an MLIL->MLIL translation
@@ -4523,11 +4541,11 @@ def _get_llil_ssa_to_mlil_expr_map(self, from_builders: bool) -> LLILSSAToMLILEx
 				# Look up the LLIL SSA expression for the old expr in its function
 				# And then store that mapping for the new function
 
+				old_llil_ssa = old_expr.function.low_level_il.ssa_form
 				old_llil_ssa_direct = old_expr.function.get_low_level_il_expr_index(old_expr.expr_index)
 				old_llil_ssa_indices = old_expr.function.get_low_level_il_expr_indexes(old_expr.expr_index)
 				for old_index in old_llil_ssa_indices:
-					old_reverse_direct = old_expr.function.low_level_il.ssa_form.get_medium_level_il_expr_index(old_index)
-					old_reverse_all = old_expr.function.low_level_il.ssa_form.get_medium_level_il_expr_indexes(old_index)
+					old_reverse_direct, old_reverse_all = reverse_lookup(old_llil_ssa, old_index)
 
 					for (new_index, new_direct) in new_indices:
 						lower_to_higher_direct = new_direct and old_reverse_direct == old_expr.expr_index
@@ -4544,13 +4562,13 @@ def _get_llil_ssa_to_mlil_expr_map(self, from_builders: bool) -> LLILSSAToMLILEx
 							higher_to_lower_direct
 						))
 		else:
+			llil_ssa = self.low_level_il.ssa_form
 			for instr in self.instructions:
 				for expr in instr.traverse(lambda e: e):
 					llil_ssa_direct = self.get_low_level_il_expr_index(expr.expr_index)
 					llil_ssa_indices = self.get_low_level_il_expr_indexes(expr.expr_index)
 					for llil_ssa_index in llil_ssa_indices:
-						reverse_direct = self.low_level_il.ssa_form.get_medium_level_il_expr_index(llil_ssa_index)
-						reverse_all = self.low_level_il.ssa_form.get_medium_level_il_expr_indexes(llil_ssa_index)
+						reverse_direct, reverse_all = reverse_lookup(llil_ssa, llil_ssa_index)
 
 						lower_to_higher_direct = reverse_direct == expr.expr_index
 						higher_to_lower_direct = llil_ssa_index == llil_ssa_direct