Move LLIL instruction retrieval into the LLIL function where it belongs

The python API was kept the same seeing as we are close to the release, will likely start deprecating some of those API's soon.

2 files changed, 62 insertions, 63 deletions

diff --git a/python/function.py b/python/function.py
index 01b482c4..eb80f9b0 100644
--- a/python/function.py
+++ b/python/function.py
@@ -1852,15 +1852,12 @@ class Function:
 			>>> func.get_low_level_il_at(func.start)
 			<il: push(rbp)>
 		"""
-		if arch is None:
-			arch = self.arch
-
-		idx = core.BNGetLowLevelILForInstruction(self.handle, arch.handle, addr)
-
 		llil = self.llil
-		if llil is None or idx == len(llil):
+		if llil is None:
+			return None
+		idx = llil.get_instruction_start(addr, arch)
+		if idx is None:
 			return None
-
 		return llil[idx]
 
 	def get_low_level_ils_at(self, addr: int,
@@ -1881,19 +1878,7 @@ class Function:
 		llil = self.llil
 		if llil is None:
 			return []
-
-		if arch is None:
-			arch = self.arch
-		count = ctypes.c_ulonglong()
-		instrs = core.BNGetLowLevelILInstructionsForAddress(self.handle, arch.handle, addr, count)
-		assert instrs is not None, "core.BNGetLowLevelILInstructionsForAddress returned None"
-		try:
-			result = []
-			for i in range(0, count.value):
-				result.append(llil[instrs[i]])
-			return result
-		finally:
-			core.BNFreeILInstructionList(instrs)
+		return [llil[i] for i in llil.get_instructions_at(addr, arch)]
 
 	def get_llil_at(self, addr: int,
 	                arch: Optional['architecture.Architecture'] = None) -> Optional['lowlevelil.LowLevelILInstruction']:
@@ -1929,33 +1914,16 @@ class Function:
 		llil = self.llil
 		if llil is None:
 			return []
-
-		if arch is None:
-			arch = self.arch
-		count = ctypes.c_ulonglong()
-		instrs = core.BNGetLowLevelILInstructionsForAddress(self.handle, arch.handle, addr, count)
-		assert instrs is not None, "core.BNGetLowLevelILInstructionsForAddress returned None"
-		try:
-			result = []
-			for i in range(0, count.value):
-				result.append(llil[instrs[i]])
-			return result
-		finally:
-			core.BNFreeILInstructionList(instrs)
+		return [llil[i] for i in llil.get_instructions_at(addr, arch)]
 
 	def get_low_level_il_exits_at(self, addr: int, arch: Optional['architecture.Architecture'] = None) -> List[int]:
-		if arch is None:
-			arch = self.arch
-		count = ctypes.c_ulonglong()
-		exits = core.BNGetLowLevelILExitsForInstruction(self.handle, arch.handle, addr, count)
-		assert exits is not None, "core.BNGetLowLevelILExitsForInstruction returned None"
-		try:
-			result = []
-			for i in range(0, count.value):
-				result.append(exits[i])
-			return result
-		finally:
-			core.BNFreeILInstructionList(exits)
+		llil = self.llil
+		if llil is None:
+			return []
+		idx = llil.get_instruction_start(addr, arch)
+		if idx is None:
+			return []
+		return llil.get_exits_for_instr(idx)
 
 	def get_constant_data(self, state: RegisterValueType, value: int, size: int = 0) -> databuffer.DataBuffer:
 		return databuffer.DataBuffer(handle=core.BNGetConstantData(self.handle, state, value, size, None))
@@ -2144,15 +2112,13 @@ class Function:
 	def get_lifted_il_at(
 	    self, addr: int, arch: Optional['architecture.Architecture'] = None
 	) -> Optional['lowlevelil.LowLevelILInstruction']:
-		if arch is None:
-			arch = self.arch
-
-		idx = core.BNGetLiftedILForInstruction(self.handle, arch.handle, addr)
-
-		if idx == len(self.lifted_il):
+		lifted_il = self.lifted_il
+		if lifted_il is None:
 			return None
-
-		return self.lifted_il[idx]
+		idx = lifted_il.get_instruction_start(addr, arch)
+		if idx is None:
+			return None
+		return lifted_il[idx]
 
 	def get_lifted_ils_at(
 	    self, addr: int, arch: Optional['architecture.Architecture'] = None
@@ -2168,16 +2134,10 @@ class Function:
 			>>> func.get_lifted_ils_at(func.start)
 			[<il: push(rbp)>]
 		"""
-		if arch is None:
-			arch = self.arch
-		count = ctypes.c_ulonglong()
-		instrs = core.BNGetLiftedILInstructionsForAddress(self.handle, arch.handle, addr, count)
-		assert instrs is not None, "core.BNGetLiftedILInstructionsForAddress returned None"
-		result = []
-		for i in range(0, count.value):
-			result.append(self.lifted_il[instrs[i]])
-		core.BNFreeILInstructionList(instrs)
-		return result
+		lifted_il = self.lifted_il
+		if lifted_il is None:
+			return []
+		return [lifted_il[i] for i in lifted_il.get_instructions_at(addr, arch)]
 
 	def get_constants_referenced_by(self, addr: int,
 	                                arch: Optional['architecture.Architecture'] = None) -> List[variable.ConstantReference]:
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 188a4623..82a09603 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -3799,6 +3799,45 @@ class LowLevelILFunction:
 			return None
 		return InstructionIndex(result)
 
+	def get_instructions_at(self, addr: int, arch: Optional['architecture.Architecture'] = None) -> List[InstructionIndex]:
+		"""
+		``get_instructions_at`` gets the InstructionIndex(s) corresponding to the given virtual address
+		See the `docs for mappings between ils <https://dev-docs.binary.ninja/dev/concepts.html#mapping-between-ils>`_ for more information.
+
+		:param int addr: virtual address of the instruction to be queried
+		:param Architecture arch: (optional) Architecture for the given function
+		:rtype: list(InstructionIndex)
+		:Example:
+
+			>>> func = next(bv.functions)
+			>>> func.llil.get_instructions_at(func.start)
+			[0]
+		"""
+		if arch is None:
+			arch = self.arch
+		count = ctypes.c_ulonglong()
+		instrs = core.BNLowLevelILGetInstructionsAt(self.handle, arch.handle, addr, count)
+		assert instrs is not None, "core.BNLowLevelILGetInstructionsAt returned None"
+		try:
+			result = []
+			for i in range(0, count.value):
+				result.append(instrs[i])
+			return result
+		finally:
+			core.BNFreeILInstructionList(instrs)
+
+	def get_exits_for_instr(self, idx: InstructionIndex) -> List[InstructionIndex]:
+		count = ctypes.c_ulonglong()
+		exits = core.BNLowLevelILGetExitsForInstruction(self.handle, idx, count)
+		assert exits is not None, "core.BNLowLevelILGetExitsForInstruction returned None"
+		try:
+			result = []
+			for i in range(0, count.value):
+				result.append(exits[i])
+			return result
+		finally:
+			core.BNFreeILInstructionList(exits)
+
 	def clear_indirect_branches(self) -> None:
 		core.BNLowLevelILClearIndirectBranches(self.handle)