From f193132145f62cac0ffeebc3ff918ebf9dd37103 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty.wagner@gmail.com>
Date: Fri, 26 Apr 2024 18:25:36 -0400
Subject: [PATCH] Add MLIL APIs for getting by instruction index, and expose
 the operation size in the Rust API

---
 rust/src/mlil/function.rs    | 14 ++++++++++++++
 rust/src/mlil/instruction.rs | 10 +++++++++-
 rust/src/mlil/lift.rs        |  1 +
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/rust/src/mlil/function.rs b/rust/src/mlil/function.rs
index 16cc51020..0b662578c 100644
--- a/rust/src/mlil/function.rs
+++ b/rust/src/mlil/function.rs
@@ -2,6 +2,7 @@ use core::hash::{Hash, Hasher};
 
 use binaryninjacore_sys::BNFreeMediumLevelILFunction;
 use binaryninjacore_sys::BNGetMediumLevelILBasicBlockList;
+use binaryninjacore_sys::BNGetMediumLevelILIndexForInstruction;
 use binaryninjacore_sys::BNGetMediumLevelILInstructionCount;
 use binaryninjacore_sys::BNGetMediumLevelILOwnerFunction;
 use binaryninjacore_sys::BNGetMediumLevelILSSAForm;
@@ -65,6 +66,19 @@ impl MediumLevelILFunction {
         self.instruction_from_idx(expr_idx).lift()
     }
 
+    pub fn instruction_from_instruction_idx(&self, instr_idx: usize) -> MediumLevelILInstruction {
+        MediumLevelILInstruction::new(self.to_owned(), unsafe {
+            BNGetMediumLevelILIndexForInstruction(self.handle, instr_idx)
+        })
+    }
+
+    pub fn lifted_instruction_from_instruction_idx(
+        &self,
+        instr_idx: usize,
+    ) -> MediumLevelILLiftedInstruction {
+        self.instruction_from_instruction_idx(instr_idx).lift()
+    }
+
     pub fn instruction_count(&self) -> usize {
         unsafe { BNGetMediumLevelILInstructionCount(self.handle) }
     }
diff --git a/rust/src/mlil/instruction.rs b/rust/src/mlil/instruction.rs
index 3b5dcfb42..bd2cc7171 100644
--- a/rust/src/mlil/instruction.rs
+++ b/rust/src/mlil/instruction.rs
@@ -18,6 +18,7 @@ pub struct MediumLevelILInstruction {
     pub function: Ref<MediumLevelILFunction>,
     pub address: u64,
     pub index: usize,
+    pub size: usize,
     pub kind: MediumLevelILInstructionKind,
 }
 
@@ -704,7 +705,12 @@ impl MediumLevelILInstruction {
             }),
             // translated directly into a list for Expression or Variables
             // TODO MLIL_MEMORY_INTRINSIC_SSA needs to be handled properly
-            MLIL_CALL_OUTPUT | MLIL_CALL_PARAM | MLIL_CALL_PARAM_SSA | MLIL_CALL_OUTPUT_SSA | MLIL_MEMORY_INTRINSIC_OUTPUT_SSA | MLIL_MEMORY_INTRINSIC_SSA => {
+            MLIL_CALL_OUTPUT
+            | MLIL_CALL_PARAM
+            | MLIL_CALL_PARAM_SSA
+            | MLIL_CALL_OUTPUT_SSA
+            | MLIL_MEMORY_INTRINSIC_OUTPUT_SSA
+            | MLIL_MEMORY_INTRINSIC_SSA => {
                 unreachable!()
             }
         };
@@ -713,6 +719,7 @@ impl MediumLevelILInstruction {
             function,
             address: op.address,
             index,
+            size: op.size,
             kind,
         }
     }
@@ -1022,6 +1029,7 @@ impl MediumLevelILInstruction {
             function: self.function.clone(),
             address: self.address,
             index: self.index,
+            size: self.size,
             kind,
         }
     }
diff --git a/rust/src/mlil/lift.rs b/rust/src/mlil/lift.rs
index 39e8e9830..e8548b064 100644
--- a/rust/src/mlil/lift.rs
+++ b/rust/src/mlil/lift.rs
@@ -27,6 +27,7 @@ pub struct MediumLevelILLiftedInstruction {
     pub function: Ref<MediumLevelILFunction>,
     pub address: u64,
     pub index: usize,
+    pub size: usize,
     pub kind: MediumLevelILLiftedInstructionKind,
 }