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.

9 files changed, 136 insertions, 164 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index c15cf24d..cb2d3632 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -11345,16 +11345,6 @@ namespace BinaryNinja {
 		*/
 		Ref<LowLevelILFunction> GetLowLevelILIfAvailable() const;
 
-		/*! Get the Low Level IL Instruction start for an instruction at an address
-
-			\param arch Architecture for the instruction
-			\param addr Address of the instruction
-			\return Start address of the instruction
-		*/
-		size_t GetLowLevelILForInstruction(Architecture* arch, uint64_t addr);
-		std::set<size_t> GetLowLevelILInstructionsForAddress(Architecture* arch, uint64_t addr);
-		std::vector<size_t> GetLowLevelILExitsForInstruction(Architecture* arch, uint64_t addr);
-
 		std::pair<DataBuffer, BNBuiltinType> GetConstantData(
 			BNRegisterValueType state, uint64_t value, size_t size = 0);
 
@@ -12289,6 +12279,9 @@ namespace BinaryNinja {
 		uint64_t GetCurrentAddress() const;
 		void SetCurrentAddress(Architecture* arch, uint64_t addr);
 		size_t GetInstructionStart(Architecture* arch, uint64_t addr);
+		std::set<size_t> GetInstructionsAt(Architecture* arch, uint64_t addr);
+
+		std::vector<size_t> GetExitsForInstruction(size_t i);
 
 		void ClearIndirectBranches();
 		void SetIndirectBranches(const std::vector<ArchAndAddr>& branches);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 3d32680c..61d7a381 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,14 +37,14 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 123
+#define BN_CURRENT_CORE_ABI_VERSION 124
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
 // will require rebuilding. The minimum version is increased when there are
 // incompatible changes that break binary compatibility, such as changes to
 // existing types or functions.
-#define BN_MINIMUM_CORE_ABI_VERSION 121
+#define BN_MINIMUM_CORE_ABI_VERSION 124
 
 #ifdef __GNUC__
 	#ifdef BINARYNINJACORE_LIBRARY
@@ -4787,11 +4787,6 @@ extern "C"
 
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelIL(BNFunction* func);
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelILIfAvailable(BNFunction* func);
-	BINARYNINJACOREAPI size_t BNGetLowLevelILForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr);
-	BINARYNINJACOREAPI size_t* BNGetLowLevelILInstructionsForAddress(
-	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
-	BINARYNINJACOREAPI size_t* BNGetLowLevelILExitsForInstruction(
-	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeILInstructionList(size_t* list);
 	BINARYNINJACOREAPI BNMediumLevelILFunction* BNGetFunctionMediumLevelIL(BNFunction* func);
 	BINARYNINJACOREAPI BNMediumLevelILFunction* BNGetFunctionMediumLevelILIfAvailable(BNFunction* func);
@@ -4837,9 +4832,6 @@ extern "C"
 
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLiftedIL(BNFunction* func);
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLiftedILIfAvailable(BNFunction* func);
-	BINARYNINJACOREAPI size_t BNGetLiftedILForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr);
-	BINARYNINJACOREAPI size_t* BNGetLiftedILInstructionsForAddress(
-	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI size_t* BNGetLiftedILFlagUsesForDefinition(
 	    BNFunction* func, size_t i, uint32_t flag, size_t* count);
 	BINARYNINJACOREAPI size_t* BNGetLiftedILFlagDefinitionsForUse(
@@ -5967,6 +5959,10 @@ extern "C"
 	BINARYNINJACOREAPI void BNLowLevelILSetCurrentSourceBlock(BNLowLevelILFunction* func, BNBasicBlock* source);
 	BINARYNINJACOREAPI size_t BNLowLevelILGetInstructionStart(
 	    BNLowLevelILFunction* func, BNArchitecture* arch, uint64_t addr);
+	BINARYNINJACOREAPI size_t* BNLowLevelILGetInstructionsAt(
+		BNLowLevelILFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI size_t* BNLowLevelILGetExitsForInstruction(
+		BNLowLevelILFunction* func, size_t instr, size_t* count);
 	BINARYNINJACOREAPI void BNLowLevelILClearIndirectBranches(BNLowLevelILFunction* func);
 	BINARYNINJACOREAPI void BNLowLevelILSetIndirectBranches(
 	    BNLowLevelILFunction* func, BNArchitectureAndAddress* branches, size_t count);
diff --git a/function.cpp b/function.cpp
index c2e52e55..1b2c0a46 100644
--- a/function.cpp
+++ b/function.cpp
@@ -460,39 +460,6 @@ Ref<LowLevelILFunction> Function::GetLowLevelILIfAvailable() const
 }
 
 
-size_t Function::GetLowLevelILForInstruction(Architecture* arch, uint64_t addr)
-{
-	return BNGetLowLevelILForInstruction(m_object, arch->GetObject(), addr);
-}
-
-
-set<size_t> Function::GetLowLevelILInstructionsForAddress(Architecture* arch, uint64_t addr)
-{
-	size_t count;
-	size_t* instrs = BNGetLowLevelILInstructionsForAddress(m_object, arch->GetObject(), addr, &count);
-
-	set<size_t> result;
-	for (size_t i = 0; i < count; i++)
-		result.insert(instrs[i]);
-
-	BNFreeILInstructionList(instrs);
-	return result;
-}
-
-
-vector<size_t> Function::GetLowLevelILExitsForInstruction(Architecture* arch, uint64_t addr)
-{
-	size_t count;
-	size_t* exits = BNGetLowLevelILExitsForInstruction(m_object, arch->GetObject(), addr, &count);
-
-	vector<size_t> result;
-	result.insert(result.end(), exits, &exits[count]);
-
-	BNFreeILInstructionList(exits);
-	return result;
-}
-
-
 RegisterValue RegisterValue::FromAPIObject(const BNRegisterValue& value)
 {
 	RegisterValue result;
@@ -795,26 +762,6 @@ Ref<LowLevelILFunction> Function::GetLiftedILIfAvailable() const
 }
 
 
-size_t Function::GetLiftedILForInstruction(Architecture* arch, uint64_t addr)
-{
-	return BNGetLiftedILForInstruction(m_object, arch->GetObject(), addr);
-}
-
-
-set<size_t> Function::GetLiftedILInstructionsForAddress(Architecture* arch, uint64_t addr)
-{
-	size_t count;
-	size_t* instrs = BNGetLiftedILInstructionsForAddress(m_object, arch->GetObject(), addr, &count);
-
-	set<size_t> result;
-	for (size_t i = 0; i < count; i++)
-		result.insert(instrs[i]);
-
-	BNFreeILInstructionList(instrs);
-	return result;
-}
-
-
 set<size_t> Function::GetLiftedILFlagUsesForDefinition(size_t i, uint32_t flag)
 {
 	size_t count;
diff --git a/lowlevelil.cpp b/lowlevelil.cpp
index 33dea219..dd3989f1 100644
--- a/lowlevelil.cpp
+++ b/lowlevelil.cpp
@@ -96,6 +96,29 @@ size_t LowLevelILFunction::GetInstructionStart(Architecture* arch, uint64_t addr
 	return BNLowLevelILGetInstructionStart(m_object, arch ? arch->GetObject() : nullptr, addr);
 }
 
+std::set<size_t> LowLevelILFunction::GetInstructionsAt(Architecture *arch, uint64_t addr)
+{
+	size_t count;
+	size_t* instructions = BNLowLevelILGetInstructionsAt(m_object, arch ? arch->GetObject() : nullptr, addr, &count);
+	std::set<size_t> result;
+	for (size_t i = 0; i < count; i++)
+		result.insert(instructions[i]);
+	BNFreeILInstructionList(instructions);
+	return result;
+}
+
+std::vector<size_t> LowLevelILFunction::GetExitsForInstruction(size_t i)
+{
+	size_t count;
+	size_t* instructions = BNLowLevelILGetExitsForInstruction(m_object, i, &count);
+	std::vector<size_t> result;
+	result.reserve(count);
+	for (size_t j = 0; j < count; j++)
+		result.push_back(instructions[j]);
+	BNFreeILInstructionList(instructions);
+	return result;
+}
+
 
 void LowLevelILFunction::ClearIndirectBranches()
 {
diff --git a/plugins/efi_resolver/src/Resolver.cpp b/plugins/efi_resolver/src/Resolver.cpp
index 5796b0e0..9699b043 100644
--- a/plugins/efi_resolver/src/Resolver.cpp
+++ b/plugins/efi_resolver/src/Resolver.cpp
@@ -308,7 +308,7 @@ vector<HighLevelILInstruction> Resolver::HighLevelILExprsAt(Ref<Function> func,
 	auto mlil = func->GetMediumLevelIL();
 	auto hlil = func->GetHighLevelIL();
 
-	size_t llilIdx = func->GetLowLevelILForInstruction(arch, addr);
+	size_t llilIdx = llil->GetInstructionStart(arch, addr);
 	size_t llilExprIdx = llil->GetIndexForInstruction(llilIdx);
 	auto mlilIdxes = llil->GetMediumLevelILExprIndexes(llilExprIdx);
 
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)
 
diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index cf3a432a..1955c147 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -93,31 +93,6 @@ where
         }
     }
 
-    /// Get all the contiguous instructions for a given location.
-    ///
-    /// NOTE: This won't get you every instruction for a location, only the instructions
-    /// that are sequential from the starting instruction.
-    pub fn instructions_at<L: Into<Location>>(&self, loc: L) -> Vec<LowLevelILInstruction<M, F>> {
-        let loc = loc.into();
-        // TODO: Instructions sharing the same address are not always sequential.
-        // Gather all of the sequential instructions with the same address and same block.
-        self.instruction_index_at(loc)
-            .map(|mut idx| {
-                let mut instructions = Vec::new();
-                let block = self.basic_block_containing_index(idx);
-                while idx.0 < self.instruction_count() {
-                    let instr = LowLevelILInstruction::new(self, idx);
-                    if instr.address() != loc.addr || instr.basic_block() != block {
-                        break;
-                    }
-                    instructions.push(instr);
-                    idx = idx.next();
-                }
-                instructions
-            })
-            .unwrap_or_default()
-    }
-
     pub fn instruction_at<L: Into<Location>>(&self, loc: L) -> Option<LowLevelILInstruction<M, F>> {
         Some(LowLevelILInstruction::new(
             self,
@@ -125,6 +100,15 @@ where
         ))
     }
 
+    /// Get all the instructions for a given location.
+    pub fn instructions_at<L: Into<Location>>(&self, loc: L) -> Vec<LowLevelILInstruction<M, F>> {
+        let loc = loc.into();
+        self.instruction_indexes_at(loc)
+            .iter()
+            .map(|idx| LowLevelILInstruction::new(self, idx))
+            .collect()
+    }
+
     pub fn instruction_index_at<L: Into<Location>>(
         &self,
         loc: L,
@@ -143,6 +127,20 @@ where
         }
     }
 
+    pub fn instruction_indexes_at<L: Into<Location>>(
+        &self,
+        loc: L,
+    ) -> Array<LowLevelInstructionIndex> {
+        let loc: Location = loc.into();
+        // If the location does not specify an architecture, use the function's architecture.
+        let arch = loc.arch.unwrap_or_else(|| self.arch());
+        let mut count = 0;
+        let indexes = unsafe {
+            BNLowLevelILGetInstructionsAt(self.handle, arch.handle, loc.addr, &mut count)
+        };
+        unsafe { Array::new(indexes, count, ()) }
+    }
+
     pub fn instruction_from_index(
         &self,
         index: LowLevelInstructionIndex,
diff --git a/rust/src/low_level_il/instruction.rs b/rust/src/low_level_il/instruction.rs
index 04a1d2d5..2d96d121 100644
--- a/rust/src/low_level_il/instruction.rs
+++ b/rust/src/low_level_il/instruction.rs
@@ -13,16 +13,16 @@
 // limitations under the License.
 
 use crate::basic_block::BasicBlock;
-use crate::rc::Ref;
+use crate::rc::{CoreArrayProvider, CoreArrayProviderInner, Ref};
 
 use super::block::LowLevelILBlock;
 use super::operation;
 use super::operation::Operation;
 use super::VisitorAction;
 use super::*;
-use binaryninjacore_sys::BNGetLowLevelILByIndex;
 use binaryninjacore_sys::BNGetLowLevelILIndexForInstruction;
 use binaryninjacore_sys::BNLowLevelILInstruction;
+use binaryninjacore_sys::{BNFreeILInstructionList, BNGetLowLevelILByIndex};
 use std::fmt::{Debug, Display, Formatter};
 
 #[repr(transparent)]
@@ -53,6 +53,22 @@ impl Display for LowLevelInstructionIndex {
     }
 }
 
+impl CoreArrayProvider for LowLevelInstructionIndex {
+    type Raw = usize;
+    type Context = ();
+    type Wrapped<'a> = Self;
+}
+
+unsafe impl CoreArrayProviderInner for LowLevelInstructionIndex {
+    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
+        BNFreeILInstructionList(raw)
+    }
+
+    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
+        Self::from(*raw)
+    }
+}
+
 // TODO: Probably want to rename this with a LowLevelIL prefix to avoid collisions when we add handlers for other ILs
 pub trait InstructionHandler<'func, M, F>
 where