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.

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,