1 files changed, 23 insertions, 25 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,