Python: Collect and pass mappings when copying MLIL functions

4 files changed, 295 insertions, 21 deletions

diff --git a/python/commonil.py b/python/commonil.py
index e7822518..ae520627 100644
--- a/python/commonil.py
+++ b/python/commonil.py
@@ -19,7 +19,7 @@
 # IN THE SOFTWARE.
 
 from dataclasses import dataclass
-from typing import Union
+from typing import Union, Optional
 from .flowgraph import FlowGraph, FlowGraphNode
 from .enums import BranchType
 from .interaction import show_graph_report
@@ -29,6 +29,9 @@ from . import mediumlevelil
 from . import highlevelil
 
 
+invalid_il_index = 0xffffffffffffffff
+
+
 # This file contains a list of top level abstract classes for implementing BNIL instructions
 @dataclass(frozen=True, repr=False, eq=False)
 class BaseILInstruction:
@@ -210,7 +213,6 @@ class AliasedVariableInstruction(VariableInstruction):
 	pass
 
 
-@dataclass
 class ILSourceLocation:
 	"""
 	ILSourceLocation is used to indicate where expressions were defined during the lifting process
@@ -220,15 +222,52 @@ class ILSourceLocation:
 	address: int
 	source_operand: int
 
+	source_llil_instruction: Optional['lowlevelil.LowLevelILInstruction'] = None
+	source_mlil_instruction: Optional['mediumlevelil.MediumLevelILInstruction'] = None
+	source_hlil_instruction: Optional['highlevelil.HighLevelILInstruction'] = None
+	il_direct: bool = True
+
+	def __init__(self, address: int, source_operand: int):
+		self.address = address
+		self.source_operand = source_operand
+
+	def __repr__(self):
+		instr = ""
+		if self.source_llil_instruction is not None:
+			instr = f" (from LLIL {self.source_llil_instruction})"
+		if self.source_mlil_instruction is not None:
+			instr = f" (from MLIL {self.source_mlil_instruction})"
+		if self.source_hlil_instruction is not None:
+			instr = f" (from HLIL {self.source_hlil_instruction})"
+		return f"<ILSourceLocation: {self.address:x}, {self.source_operand}{instr}>"
+
+	def __hash__(self):
+		return hash((self.address, self.source_operand))
+
+	def __eq__(self, other):
+		if not isinstance(other, ILSourceLocation):
+			return False
+		return self.address == other.address and self.source_operand == other.source_operand
+
 	@classmethod
 	def from_instruction(
 			cls,
-			instr: Union['lowlevelil.LowLevelILInstruction', 'mediumlevelil.MediumLevelILInstruction', 'highlevelil.HighLevelILInstruction']
+			instr: Union['lowlevelil.LowLevelILInstruction', 'mediumlevelil.MediumLevelILInstruction', 'highlevelil.HighLevelILInstruction'],
+			il_direct: bool = True
 	) -> 'ILSourceLocation':
 		"""
 		Get the source location of a given instruction
 		:param instr: Instruction, Low, Medium, or High level
 		:return: Its location
 		"""
-		return cls(instr.address, instr.source_operand)
-
+		loc = cls(instr.address, instr.source_operand)
+		if isinstance(instr, lowlevelil.LowLevelILInstruction):
+			loc.source_llil_instruction = instr
+		elif isinstance(instr, mediumlevelil.MediumLevelILInstruction):
+			loc.source_mlil_instruction = instr
+		elif isinstance(instr, highlevelil.HighLevelILInstruction):
+			loc.source_hlil_instruction = instr
+		else:
+			log_warn(f"Unknown instruction type {type(instr)}")
+		loc.il_direct = il_direct
+		return loc
diff --git a/python/examples/wf_test_copy_expr.py b/python/examples/wf_test_copy_expr.py
index 7a3da06c..77e6e159 100644
--- a/python/examples/wf_test_copy_expr.py
+++ b/python/examples/wf_test_copy_expr.py
@@ -1,3 +1,4 @@
+import functools
 import json
 import math
 
@@ -67,7 +68,14 @@ def assert_llil_eq(old_insn: LowLevelILInstruction, new_insn: LowLevelILInstruct
             assert old_op[1] == new_op[1], err_msg
 
 
-def assert_mlil_eq(old_insn: LowLevelILInstruction, new_insn: LowLevelILInstruction):
+@functools.lru_cache(maxsize=8)
+def get_mlil_maps(mlil: MediumLevelILFunction, builders: bool) -> Tuple[LLILSSAToMLILInstructionMapping, LLILSSAToMLILExpressionMapping]:
+    instr_map = mlil._get_llil_ssa_to_mlil_instr_map(builders)
+    expr_map = mlil._get_llil_ssa_to_mlil_expr_map(builders)
+    return instr_map, expr_map
+
+
+def assert_mlil_eq(old_insn: MediumLevelILInstruction, new_insn: MediumLevelILInstruction):
     """
     Make sure that these two instructions are the same (probably correct). Asserts otherwise.
 
@@ -78,10 +86,23 @@ def assert_mlil_eq(old_insn: LowLevelILInstruction, new_insn: LowLevelILInstruct
     """
     err_msg = (hex(old_insn.address), old_insn, new_insn)
     assert old_insn.operation == new_insn.operation, err_msg
-    # assert old_insn.attributes == new_insn.attributes, err_msg
+    assert old_insn.attributes == new_insn.attributes, err_msg
     assert old_insn.size == new_insn.size, err_msg
     assert old_insn.source_location == new_insn.source_location, err_msg
     assert len(old_insn.operands) == len(new_insn.operands), err_msg
+    # Type only applies once we've generated SSA form (probably not consistent)
+    # assert old_insn.expr_type == new_insn.expr_type, f"{err_msg} {old_insn.expr_type} {new_insn.expr_type}"
+
+    instr_map, expr_map = get_mlil_maps(new_insn.function, True)
+
+    # Compare that the instruction's LLIL SSA map is the same as the old function
+    if old_insn.instr_index is not None and old_insn.function.get_expr_index_for_instruction(old_insn.instr_index) == old_insn.expr_index:
+        old_llil_ssa = old_insn.function.get_low_level_il_instruction_index(old_insn.instr_index)
+        if old_llil_ssa is not None:
+            assert [mlil for (llil, mlil) in instr_map.items() if llil == old_llil_ssa] == [new_insn.instr_index], err_msg
+        else:
+            assert [mlil for (llil, mlil) in instr_map.items() if llil == old_llil_ssa] == [], err_msg
+
     # Can't compare operands directly since IL expression indices might change when
     # copying an instruction to another function
     for i, (old_op, new_op) in enumerate(zip(old_insn.detailed_operands, new_insn.detailed_operands)):
@@ -231,6 +252,11 @@ def mlil_action(context: AnalysisContext):
         report.append(FlowGraphReport("new graph", new_mlil.create_graph_immediate(settings)))
         show_report_collection("copy expr test", report)
 
+    # Check expr mappings are the same
+    new_map = list(sorted(new_mlil._get_llil_ssa_to_mlil_expr_map(True), key=lambda o: (o.lower_index, o.higher_index)))
+    old_map = list(sorted(old_mlil._get_llil_ssa_to_mlil_expr_map(False), key=lambda o: (o.lower_index, o.higher_index)))
+    assert old_map == new_map
+
     # Check all BBs have all the same instructions
     # Technically, this misses any instructions outside a BB, but those are not
     # picked up by analysis anyway, and therefore don't matter.
@@ -240,15 +266,26 @@ def mlil_action(context: AnalysisContext):
         for old_insn, new_insn in zip(old_bb, new_bb):
             assert_mlil_eq(old_insn, new_insn)
 
+    # Make sure mappings update correctly following set
+    new_map = list(sorted(new_mlil._get_llil_ssa_to_mlil_expr_map(True), key=lambda o: (o.lower_index, o.higher_index)))
+    context.mlil = new_mlil
+    newer_map = list(sorted(context.mlil._get_llil_ssa_to_mlil_expr_map(False), key=lambda o: (o.lower_index, o.higher_index)))
+    assert new_map == newer_map
+
 
-wf = Workflow("core.function.metaAnalysis").clone("TestCopyExpr")
+wf = Workflow("core.function.metaAnalysis").clone("core.function.metaAnalysis")
 
 # Define the custom activity configuration
 wf.register_activity(Activity(
     configuration=json.dumps({
         "name": "extension.test_copy_expr.lil_action",
         "title": "Lifted IL copy_expr Test",
-        "description": "Makes sure copy_expr works on Lifted IL functions."
+        "description": "Makes sure copy_expr works on Lifted IL functions.",
+        "eligibility": {
+            "auto": {
+                "default": False
+            }
+        }
     }),
     action=lil_action
 ))
@@ -256,7 +293,12 @@ wf.register_activity(Activity(
     configuration=json.dumps({
         "name": "extension.test_copy_expr.llil_action",
         "title": "Low Level IL copy_expr Test",
-        "description": "Makes sure copy_expr works on Low Level IL functions."
+        "description": "Makes sure copy_expr works on Low Level IL functions.",
+        "eligibility": {
+            "auto": {
+                "default": False
+            }
+        }
     }),
     action=llil_action
 ))
@@ -264,7 +306,12 @@ wf.register_activity(Activity(
     configuration=json.dumps({
         "name": "extension.test_copy_expr.mlil_action",
         "title": "Medium Level IL copy_expr Test",
-        "description": "Makes sure copy_expr works on Medium Level IL functions."
+        "description": "Makes sure copy_expr works on Medium Level IL functions.",
+        "eligibility": {
+            "auto": {
+                "default": False
+            }
+        }
     }),
     action=mlil_action
 ))
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index cc99db8f..5d386b3f 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -20,7 +20,7 @@
 
 import ctypes
 import struct
-from typing import (Optional, List, Union, Mapping,
+from typing import (Optional, List, Union, Mapping, MutableMapping,
 	Generator, NewType, Tuple, ClassVar, Dict, Set, Callable, Any, Iterator, overload)
 from dataclasses import dataclass
 from . import deprecation
@@ -43,7 +43,7 @@ from .commonil import (
     BaseILInstruction, Constant, BinaryOperation, UnaryOperation, Comparison, SSA, Phi, FloatingPoint, ControlFlow,
     Terminal, Call, Localcall, Syscall, Tailcall, Return, Signed, Arithmetic, Carry, DoublePrecision, Memory, Load,
     Store, RegisterStack, SetVar, Intrinsic, VariableInstruction, SSAVariableInstruction, AliasedVariableInstruction,
-    ILSourceLocation
+    ILSourceLocation, invalid_il_index
 )
 
 TokenList = List['function.InstructionTextToken']
@@ -61,6 +61,30 @@ MediumLevelILOperandType = Union[int, float, 'MediumLevelILOperationAndSize', 'M
 MediumLevelILVisitorCallback = Callable[[str, MediumLevelILOperandType, str, Optional['MediumLevelILInstruction']], bool]
 StringOrType = Union[str, '_types.Type', '_types.TypeBuilder']
 ILInstructionAttributeSet = Union[Set[ILInstructionAttribute], List[ILInstructionAttribute]]
+LLILSSAToMLILInstructionMapping = MutableMapping['lowlevelil.InstructionIndex', InstructionIndex]
+
+
+@dataclass(frozen=True)
+class LLILSSAToMLILExpressionMap:
+	lower_index: 'lowlevelil.ExpressionIndex'
+	higher_index: ExpressionIndex
+	map_lower_to_higher: bool
+	map_higher_to_lower: bool
+	lower_to_higher_direct: bool
+	higher_to_lower_direct: bool
+
+	def _to_core_struct(self):
+		result = core.BNExprMapInfo()
+		result.lowerIndex = self.lower_index
+		result.higherIndex = self.higher_index
+		result.mapLowerToHigher = self.map_lower_to_higher
+		result.mapHigherToLower = self.map_higher_to_lower
+		result.lowerToHigherDirect = self.lower_to_higher_direct
+		result.higherToLowerDirect = self.higher_to_lower_direct
+		return result
+
+
+LLILSSAToMLILExpressionMapping = List['LLILSSAToMLILExpressionMap']
 
 
 @dataclass(frozen=True, repr=False, order=True)
@@ -3300,6 +3324,11 @@ class MediumLevelILFunction:
 		self._arch = _arch
 		self._source_function = _source_function
 
+		self._mlil_to_mlil_expr_map: dict['MediumLevelILInstruction', List[Tuple[ExpressionIndex, bool]]] = {}
+		self._mlil_to_mlil_instr_map: dict['MediumLevelILInstruction', List[Tuple[InstructionIndex, bool]]] = {}
+		self._llil_ssa_to_mlil_expr_map: dict['lowlevelil.LowLevelILInstruction', List[Tuple[ExpressionIndex, bool]]] = {}
+		self._llil_ssa_to_mlil_instr_map: dict['lowlevelil.LowLevelILInstruction', List[Tuple[InstructionIndex, bool]]] = {}
+
 	def __del__(self):
 		if core is not None:
 			core.BNFreeMediumLevelILFunction(self.handle)
@@ -3546,7 +3575,7 @@ class MediumLevelILFunction:
 		elif isinstance(operation, MediumLevelILOperation):
 			_operation = operation.value
 		if source_location is not None:
-			return ExpressionIndex(core.BNMediumLevelILAddExprWithLocation(
+			index = ExpressionIndex(core.BNMediumLevelILAddExprWithLocation(
 				self.handle,
 				_operation,
 				source_location.address,
@@ -3558,6 +3587,8 @@ class MediumLevelILFunction:
 				d,
 				e
 			))
+			self._record_mlil_to_mlil_expr_map(index, source_location)
+			return index
 		else:
 			return ExpressionIndex(core.BNMediumLevelILAddExpr(self.handle, _operation, size, a, b, c, d, e))
 
@@ -3914,7 +3945,6 @@ class MediumLevelILFunction:
 		new_index = do_copy(expr, dest, sub_expr_handler)
 		# Copy expression metadata as well
 		dest.set_expr_attributes(new_index, expr.attributes)
-		dest.set_expr_type(new_index, self.get_expr_type(expr.expr_index))
 		return new_index
 
 	def translate(
@@ -3944,7 +3974,7 @@ class MediumLevelILFunction:
 			for instr_index in range(block.start, block.end):
 				instr: MediumLevelILInstruction = self[InstructionIndex(instr_index)]
 				propagated_func.set_current_address(instr.address, block.arch)
-				propagated_func.append(expr_handler(propagated_func, block, instr))
+				propagated_func.append(expr_handler(propagated_func, block, instr), ILSourceLocation.from_instruction(instr))
 
 		return propagated_func
 
@@ -3968,7 +3998,7 @@ class MediumLevelILFunction:
 			result |= flag.value
 		core.BNSetMediumLevelILExprAttributes(self.handle, expr, result)
 
-	def append(self, expr: ExpressionIndex) -> InstructionIndex:
+	def append(self, expr: ExpressionIndex, source_location: Optional['ILSourceLocation'] = None) -> InstructionIndex:
 		"""
 		``append`` adds the ExpressionIndex ``expr`` to the current MediumLevelILFunction.
 
@@ -3976,7 +4006,115 @@ class MediumLevelILFunction:
 		:return: Index of added instruction in the current function
 		:rtype: int
 		"""
-		return InstructionIndex(core.BNMediumLevelILAddInstruction(self.handle, expr))
+		index = InstructionIndex(core.BNMediumLevelILAddInstruction(self.handle, expr))
+		self._record_mlil_to_mlil_instr_map(index, source_location)
+		return index
+
+	def _record_mlil_to_mlil_instr_map(self, index, source_location: 'ILSourceLocation'):
+		# Update internal mappings to remember this
+		if source_location is not None:
+			if source_location.source_mlil_instruction is not None:
+				if source_location.source_mlil_instruction not in self._mlil_to_mlil_instr_map:
+					self._mlil_to_mlil_instr_map[source_location.source_mlil_instruction] = []
+				self._mlil_to_mlil_instr_map[source_location.source_mlil_instruction].append((index, source_location.il_direct))
+			if source_location.source_llil_instruction is not None \
+				and source_location.source_llil_instruction.function.il_form == FunctionGraphType.LowLevelILSSAFormFunctionGraph:
+				if source_location.source_llil_instruction not in self._llil_ssa_to_mlil_instr_map:
+					self._llil_ssa_to_mlil_instr_map[source_location.source_llil_instruction] = []
+				self._llil_ssa_to_mlil_instr_map[source_location.source_llil_instruction].append((index, source_location.il_direct))
+
+	def _record_mlil_to_mlil_expr_map(self, index, source_location: 'ILSourceLocation'):
+		# Update internal mappings to remember this
+		if source_location.source_mlil_instruction is not None:
+			if source_location.source_mlil_instruction not in self._mlil_to_mlil_expr_map:
+				self._mlil_to_mlil_expr_map[source_location.source_mlil_instruction] = []
+			self._mlil_to_mlil_expr_map[source_location.source_mlil_instruction].append((index, source_location.il_direct))
+		if source_location.source_llil_instruction is not None \
+			and source_location.source_llil_instruction.function.il_form == FunctionGraphType.LowLevelILSSAFormFunctionGraph:
+			if source_location.source_llil_instruction not in self._llil_ssa_to_mlil_expr_map:
+				self._llil_ssa_to_mlil_expr_map[source_location.source_llil_instruction] = []
+			self._llil_ssa_to_mlil_expr_map[source_location.source_llil_instruction].append((index, source_location.il_direct))
+
+	def _get_llil_ssa_to_mlil_instr_map(self, from_builders: bool) -> LLILSSAToMLILInstructionMapping:
+		llil_ssa_to_mlil_instr_map = {}
+
+		if from_builders:
+			for (old_instr, new_indices) in self._mlil_to_mlil_instr_map.items():
+				old_instr: MediumLevelILInstruction
+				new_indices: List[InstructionIndex]
+
+				# Look up the LLIL SSA instruction for the old instr in its function
+				# And then store that mapping for the new function
+
+				for (new_index, new_direct) in new_indices:
+					# Instructions are always mapped 1 to 1. If the map is marked indirect
+					# then just ignore it
+					if new_direct:
+						old_llil_ssa_index = old_instr.function.get_low_level_il_instruction_index(old_instr.instr_index)
+						if old_llil_ssa_index is not None:
+							llil_ssa_to_mlil_instr_map[old_llil_ssa_index] = new_index
+		else:
+			for instr in self.instructions:
+				llil_ssa_index = self.get_low_level_il_instruction_index(instr.instr_index)
+				llil_ssa_to_mlil_instr_map[llil_ssa_index] = instr.instr_index
+
+		return llil_ssa_to_mlil_instr_map
+
+	def _get_llil_ssa_to_mlil_expr_map(self, from_builders: bool) -> LLILSSAToMLILExpressionMapping:
+		llil_ssa_to_mlil_expr_map = []
+
+		if from_builders:
+			for (old_expr, new_indices) in self._mlil_to_mlil_expr_map.items():
+				old_expr: MediumLevelILInstruction
+				new_indices: List[ExpressionIndex]
+
+				# 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_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)
+
+					for (new_index, new_direct) in new_indices:
+						lower_to_higher_direct = new_direct and old_reverse_direct == old_expr.expr_index
+						higher_to_lower_direct = new_direct and old_index == old_llil_ssa_direct
+						map_lower_to_higher = old_expr.expr_index in old_reverse_all
+						map_higher_to_lower = True
+
+						llil_ssa_to_mlil_expr_map.append(LLILSSAToMLILExpressionMap(
+							old_index,
+							new_index,
+							map_lower_to_higher,
+							map_higher_to_lower,
+							lower_to_higher_direct,
+							higher_to_lower_direct
+						))
+		else:
+			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)
+
+						lower_to_higher_direct = reverse_direct == expr.expr_index
+						higher_to_lower_direct = llil_ssa_index == llil_ssa_direct
+						map_lower_to_higher = expr.expr_index in reverse_all
+						map_higher_to_lower = True
+
+						llil_ssa_to_mlil_expr_map.append(LLILSSAToMLILExpressionMap(
+							llil_ssa_index,
+							expr.expr_index,
+							map_lower_to_higher,
+							map_higher_to_lower,
+							lower_to_higher_direct,
+							higher_to_lower_direct
+						))
+
+		return llil_ssa_to_mlil_expr_map
 
 	def nop(self, loc: Optional['ILSourceLocation'] = None) -> ExpressionIndex:
 		"""
@@ -5619,7 +5757,9 @@ class MediumLevelILFunction:
 		:rtype: ExpressionIndex
 		"""
 		if loc is not None:
-			return ExpressionIndex(core.BNMediumLevelILGotoWithLocation(self.handle, label.handle, loc.address, loc.source_operand))
+			index = ExpressionIndex(core.BNMediumLevelILGotoWithLocation(self.handle, label.handle, loc.address, loc.source_operand))
+			self._record_mlil_to_mlil_expr_map(index, loc)
+			return index
 		else:
 			return ExpressionIndex(core.BNMediumLevelILGoto(self.handle, label.handle))
 
@@ -5639,7 +5779,9 @@ class MediumLevelILFunction:
 		:rtype: ExpressionIndex
 		"""
 		if loc is not None:
-			return ExpressionIndex(core.BNMediumLevelILIfWithLocation(self.handle, operand, t.handle, f.handle, loc.address, loc.source_operand))
+			index = ExpressionIndex(core.BNMediumLevelILIfWithLocation(self.handle, operand, t.handle, f.handle, loc.address, loc.source_operand))
+			self._record_mlil_to_mlil_expr_map(index, loc)
+			return index
 		else:
 			return ExpressionIndex(core.BNMediumLevelILIf(self.handle, operand, t.handle, f.handle))
 
diff --git a/python/workflow.py b/python/workflow.py
index c63c20b9..bf48e5d8 100644
--- a/python/workflow.py
+++ b/python/workflow.py
@@ -112,7 +112,53 @@ class AnalysisContext:
 
 	@mlil.setter
 	def mlil(self, value: mediumlevelil.MediumLevelILFunction) -> None:
-		core.BNSetMediumLevelILFunction(self.handle, value.handle)
+		self.set_mlil_function(value)
+
+	def set_mlil_function(
+		self,
+		new_func: mediumlevelil.MediumLevelILFunction,
+		llil_ssa_to_mlil_instr_map: Optional['mediumlevelil.LLILSSAToMLILInstructionMapping'] = None,
+		llil_ssa_to_mlil_expr_map: Optional['mediumlevelil.LLILSSAToMLILExpressionMapping'] = None,
+	) -> None:
+		"""
+		Set the Medium Level IL function in the current analysis, giving updated
+		Low Level IL (SSA) to Medium Level IL instruction and expression mappings.
+		:param new_func: New MLIL function
+		:param llil_ssa_to_mlil_instr_map: Mapping from every LLIL SSA instruction to
+		                                   every MLIL instruction
+		:param llil_ssa_to_mlil_expr_map: Mapping from every LLIL SSA expression to
+		                                  one or more MLIL expressions (first expression
+		                                  will be the primary)
+		"""
+		if llil_ssa_to_mlil_instr_map is None or llil_ssa_to_mlil_expr_map is None:
+			# Build up maps from existing data in the function
+			llil_ssa_to_mlil_instr_map = new_func._get_llil_ssa_to_mlil_instr_map(True)
+			llil_ssa_to_mlil_expr_map = new_func._get_llil_ssa_to_mlil_expr_map(True)
+
+		# Number of instructions
+		instr_count = 0
+		if len(llil_ssa_to_mlil_instr_map) > 0:
+			instr_count = max(llil_ssa_to_mlil_instr_map.keys()) + 1
+		ffi_instr_map = (ctypes.c_size_t * instr_count)()
+		for i in range(instr_count):
+			ffi_instr_map[i] = 0xffffffffffffffff
+		for (key, value) in llil_ssa_to_mlil_instr_map.items():
+			ffi_instr_map[key] = value
+
+		# Number of map entries, not highest index
+		expr_count = len(llil_ssa_to_mlil_expr_map)
+		ffi_expr_map = (core.BNExprMapInfo * expr_count)()
+		for i, map in enumerate(llil_ssa_to_mlil_expr_map):
+			ffi_expr_map[i] = map._to_core_struct()
+
+		core.BNSetMediumLevelILFunction(
+			self.handle,
+			new_func.handle,
+			ffi_instr_map,
+			instr_count,
+			ffi_expr_map,
+			expr_count
+		)
 
 	@property
 	def hlil(self) -> Optional[highlevelil.HighLevelILFunction]: