From e39b96fe08b3462d3cc6f2d68e9246fbcb1fd462 Mon Sep 17 00:00:00 2001
From: Mark Rowe <mrowe@bdash.net.nz>
Date: Tue, 27 May 2025 12:34:30 -0700
Subject: [PATCH] [Rust] Add LowLevelILInstruction::basic_block /
 LowLevelILFunction::basic_block_containing_index

---
 rust/src/low_level_il/function.rs    | 15 +++++++++++++++
 rust/src/low_level_il/instruction.rs | 10 ++++++++++
 2 files changed, 25 insertions(+)

diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index 495aae7905..039a523c43 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -159,6 +159,21 @@ where
             Array::new(blocks, count, context)
         }
     }
+
+    /// Returns the [`BasicBlock`] at the given instruction `index`.
+    ///
+    /// You can also retrieve this using [`LowLevelILInstruction::basic_block`].
+    pub fn basic_block_containing_index(
+        &self,
+        index: LowLevelInstructionIndex,
+    ) -> Option<Ref<BasicBlock<LowLevelILBlock<M, F>>>> {
+        let block = unsafe { BNGetLowLevelILBasicBlockForInstruction(self.handle, index.0) };
+        if block.is_null() {
+            None
+        } else {
+            Some(unsafe { BasicBlock::ref_from_raw(block, LowLevelILBlock { function: self }) })
+        }
+    }
 }
 
 impl<M: FunctionMutability> LowLevelILFunction<M, NonSSA> {
diff --git a/rust/src/low_level_il/instruction.rs b/rust/src/low_level_il/instruction.rs
index ef546412ed..2ba57523a4 100644
--- a/rust/src/low_level_il/instruction.rs
+++ b/rust/src/low_level_il/instruction.rs
@@ -12,6 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::basic_block::BasicBlock;
+use crate::rc::Ref;
+
+use super::block::LowLevelILBlock;
 use super::operation;
 use super::operation::Operation;
 use super::VisitorAction;
@@ -98,6 +102,12 @@ where
     pub fn into_raw(&self) -> BNLowLevelILInstruction {
         unsafe { BNGetLowLevelILByIndex(self.function.handle, self.expr_idx().0) }
     }
+
+    /// Returns the [`BasicBlock`] containing the given [`LowLevelILInstruction`].
+    pub fn basic_block(&self) -> Option<Ref<BasicBlock<LowLevelILBlock<'func, M, F>>>> {
+        // TODO: We might be able to .expect this if we guarantee that self.index is valid.
+        self.function.basic_block_containing_index(self.index)
+    }
 }
 
 impl<M, F> Debug for LowLevelILInstruction<'_, M, F>