[Rust] Retrieve architecture if missing from location passed to `MediumLevelILFunction::instruction_from_index`

Fixes https://github.com/Vector35/binaryninja-api/issues/6876

3 files changed, 12 insertions, 7 deletions

diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index 8af89652..2eaf7f11 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -106,7 +106,8 @@ where
     ) -> Option<LowLevelInstructionIndex> {
         use binaryninjacore_sys::BNLowLevelILGetInstructionStart;
         let loc: Location = loc.into();
-        let arch = loc.arch.unwrap_or_else(|| *self.arch().as_ref());
+        // If the location does not specify an architecture, use the function's architecture.
+        let arch = loc.arch.unwrap_or_else(|| self.arch());
         let instr_idx =
             unsafe { BNLowLevelILGetInstructionStart(self.handle, arch.handle, loc.addr) };
         // `instr_idx` will equal self.instruction_count() if the instruction is not valid.
diff --git a/rust/src/medium_level_il/function.rs b/rust/src/medium_level_il/function.rs
index 4206f7dd..54df3d12 100644
--- a/rust/src/medium_level_il/function.rs
+++ b/rust/src/medium_level_il/function.rs
@@ -47,11 +47,10 @@ impl MediumLevelILFunction {
         loc: L,
     ) -> Option<MediumLevelInstructionIndex> {
         let loc: Location = loc.into();
-        let arch = loc
-            .arch
-            .map(|a| a.handle)
-            .unwrap_or_else(std::ptr::null_mut);
-        let instr_idx = unsafe { BNMediumLevelILGetInstructionStart(self.handle, arch, loc.addr) };
+        // If the location does not specify an architecture, use the function's architecture.
+        let arch = loc.arch.unwrap_or_else(|| self.function().arch());
+        let instr_idx =
+            unsafe { BNMediumLevelILGetInstructionStart(self.handle, arch.handle, loc.addr) };
         // `instr_idx` will equal self.instruction_count() if the instruction is not valid.
         if instr_idx >= self.instruction_count() {
             None
diff --git a/rust/tests/medium_level_il.rs b/rust/tests/medium_level_il.rs
index de79a47d..a921965f 100644
--- a/rust/tests/medium_level_il.rs
+++ b/rust/tests/medium_level_il.rs
@@ -164,9 +164,14 @@ fn test_mlil_possible_values() {
         _ => panic!("Expected SetVar"),
     }
 
+    // As an aside, we also test to make sure `instruction_index_at` works.
+    let instr_21_idx = mlil_function
+        .instruction_index_at(image_base + 0x0002af9a)
+        .expect("Failed to get instruction index");
+
     // 21 @ 0002af9a  (MLIL_RET return (MLIL_VAR.d eax_2))
     let instr_21 = mlil_function
-        .instruction_from_index(MediumLevelInstructionIndex(21))
+        .instruction_from_index(instr_21_idx)
         .expect("Failed to get instruction");
     let lifted_instr_21 = instr_21.lift();
     match lifted_instr_21.kind {