From bec653143bd78f1662ed1a556c4b00471f395fcc Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty.wagner@gmail.com>
Date: Fri, 26 Apr 2024 18:26:12 -0400
Subject: [PATCH] Add HLIL APIs to fetch the root or by instruction index, and
 expose the operation size in the Rust API

---
 rust/src/hlil/function.rs    | 41 ++++++++++++++++++++++++++++++++++++
 rust/src/hlil/instruction.rs |  3 +++
 rust/src/hlil/lift.rs        |  1 +
 rust/src/hlil/operation.rs   |  2 +-
 4 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/rust/src/hlil/function.rs b/rust/src/hlil/function.rs
index 4bad7f0f0..25608d714 100644
--- a/rust/src/hlil/function.rs
+++ b/rust/src/hlil/function.rs
@@ -2,8 +2,10 @@ use std::hash::{Hash, Hasher};
 
 use binaryninjacore_sys::BNFreeHighLevelILFunction;
 use binaryninjacore_sys::BNGetHighLevelILBasicBlockList;
+use binaryninjacore_sys::BNGetHighLevelILIndexForInstruction;
 use binaryninjacore_sys::BNGetHighLevelILInstructionCount;
 use binaryninjacore_sys::BNGetHighLevelILOwnerFunction;
+use binaryninjacore_sys::BNGetHighLevelILRootExpr;
 use binaryninjacore_sys::BNGetHighLevelILSSAForm;
 use binaryninjacore_sys::BNHighLevelILFunction;
 use binaryninjacore_sys::BNNewHighLevelILFunctionReference;
@@ -52,6 +54,29 @@ impl HighLevelILFunction {
         self.instruction_from_idx(expr_idx).lift()
     }
 
+    pub fn instruction_from_instruction_idx(&self, instr_idx: usize) -> HighLevelILInstruction {
+        HighLevelILInstruction::new(self.as_non_ast(), unsafe {
+            BNGetHighLevelILIndexForInstruction(self.handle, instr_idx)
+        })
+    }
+
+    pub fn lifted_instruction_from_instruction_idx(
+        &self,
+        instr_idx: usize,
+    ) -> HighLevelILLiftedInstruction {
+        self.instruction_from_instruction_idx(instr_idx).lift()
+    }
+
+    pub fn root(&self) -> HighLevelILInstruction {
+        HighLevelILInstruction::new(self.as_ast(), unsafe {
+            BNGetHighLevelILRootExpr(self.handle)
+        })
+    }
+
+    pub fn lifted_root(&self) -> HighLevelILLiftedInstruction {
+        self.root().lift()
+    }
+
     pub fn instruction_count(&self) -> usize {
         unsafe { BNGetHighLevelILInstructionCount(self.handle) }
     }
@@ -81,6 +106,22 @@ impl HighLevelILFunction {
 
         unsafe { Array::new(blocks, count, context) }
     }
+
+    pub fn as_ast(&self) -> Ref<HighLevelILFunction> {
+        Self {
+            handle: self.handle,
+            full_ast: true,
+        }
+        .to_owned()
+    }
+
+    pub fn as_non_ast(&self) -> Ref<HighLevelILFunction> {
+        Self {
+            handle: self.handle,
+            full_ast: false,
+        }
+        .to_owned()
+    }
 }
 
 impl ToOwned for HighLevelILFunction {
diff --git a/rust/src/hlil/instruction.rs b/rust/src/hlil/instruction.rs
index 9bffadaf4..7e77e379e 100644
--- a/rust/src/hlil/instruction.rs
+++ b/rust/src/hlil/instruction.rs
@@ -16,6 +16,7 @@ pub struct HighLevelILInstruction {
     pub function: Ref<HighLevelILFunction>,
     pub address: u64,
     pub index: usize,
+    pub size: usize,
     pub kind: HighLevelILInstructionKind,
 }
 
@@ -629,6 +630,7 @@ impl HighLevelILInstruction {
             function,
             address: op.address,
             index,
+            size: op.size,
             kind,
         }
     }
@@ -878,6 +880,7 @@ impl HighLevelILInstruction {
             function: self.function.clone(),
             address: self.address,
             index: self.index,
+            size: self.size,
             kind,
         }
     }
diff --git a/rust/src/hlil/lift.rs b/rust/src/hlil/lift.rs
index 74ae0c64e..731a785c1 100644
--- a/rust/src/hlil/lift.rs
+++ b/rust/src/hlil/lift.rs
@@ -24,6 +24,7 @@ pub struct HighLevelILLiftedInstruction {
     pub function: Ref<HighLevelILFunction>,
     pub address: u64,
     pub index: usize,
+    pub size: usize,
     pub kind: HighLevelILLiftedInstructionKind,
 }
 
diff --git a/rust/src/hlil/operation.rs b/rust/src/hlil/operation.rs
index 965d95173..ee0d437b5 100644
--- a/rust/src/hlil/operation.rs
+++ b/rust/src/hlil/operation.rs
@@ -9,7 +9,7 @@ use super::HighLevelILLiftedInstruction;
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct GotoLabel {
     pub(crate) function: Ref<Function>,
-    pub(crate) target: u64,
+    pub target: u64,
 }
 
 impl GotoLabel {