2 files changed, 41 insertions, 27 deletions

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