Rust APIs for custom function lifters

2 files changed, 46 insertions, 3 deletions

diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index 32e1f921..a0db838b 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -209,6 +209,20 @@ where
             Some(unsafe { BasicBlock::ref_from_raw(block, LowLevelILBlock { function: self }) })
         }
     }
+
+    pub fn set_indirect_branches(&self, branches: &[Location]) {
+        let mut bn_branches: Box<[BNArchitectureAndAddress]> = branches
+            .iter()
+            .map(|loc| BNArchitectureAndAddress {
+                address: loc.addr,
+                arch: loc.arch.unwrap_or_else(|| self.arch()).handle,
+            })
+            .collect();
+
+        unsafe {
+            BNLowLevelILSetIndirectBranches(self.handle, bn_branches.as_mut_ptr(), branches.len());
+        }
+    }
 }
 
 impl<M: FunctionMutability> LowLevelILFunction<M, NonSSA> {
diff --git a/rust/src/low_level_il/lifting.rs b/rust/src/low_level_il/lifting.rs
index d162b3f6..202c9ed6 100644
--- a/rust/src/low_level_il/lifting.rs
+++ b/rust/src/low_level_il/lifting.rs
@@ -14,8 +14,10 @@
 
 use std::marker::PhantomData;
 
-use binaryninjacore_sys::{BNAddLowLevelILLabelForAddress, BNLowLevelILOperation};
-use binaryninjacore_sys::{BNLowLevelILLabel, BNRegisterOrConstant};
+use binaryninjacore_sys::{
+    BNAddLowLevelILLabelForAddress, BNLowLevelILClearIndirectBranches, BNLowLevelILLabel,
+    BNLowLevelILOperation, BNRegisterOrConstant, BNSetLowLevelILExprAttributes,
+};
 
 use super::*;
 use crate::architecture::{Architecture, FlagWriteId, RegisterId};
@@ -23,7 +25,8 @@ use crate::architecture::{CoreRegister, Register as ArchReg};
 use crate::architecture::{
     Flag, FlagClass, FlagCondition, FlagGroup, FlagRole, FlagWrite, Intrinsic,
 };
-use crate::function::Location;
+use crate::basic_block::BasicBlock;
+use crate::function::{Location, NativeBlock};
 
 pub trait LiftableLowLevelIL<'func> {
     type Result: ExpressionResultType;
@@ -1512,6 +1515,13 @@ impl LowLevelILMutableFunction {
         }
     }
 
+    pub fn set_current_source_block(&self, source: &BasicBlock<NativeBlock>) {
+        use binaryninjacore_sys::BNLowLevelILSetCurrentSourceBlock;
+        unsafe {
+            BNLowLevelILSetCurrentSourceBlock(self.handle, source.handle);
+        }
+    }
+
     pub fn label_for_address<L: Into<Location>>(&self, loc: L) -> Option<LowLevelILLabel> {
         use binaryninjacore_sys::BNGetLowLevelILLabelForAddress;
 
@@ -1561,6 +1571,25 @@ impl LowLevelILMutableFunction {
         }
         *label = new_label;
     }
+
+    pub fn set_expr_attributes(
+        &self,
+        expr: LowLevelExpressionIndex,
+        value: &ILInstructionAttributeSet,
+    ) {
+        let mut result = 0u32;
+        for flag in value {
+            result |= flag.value();
+        }
+
+        unsafe {
+            BNSetLowLevelILExprAttributes(self.handle, expr.0, result);
+        }
+    }
+
+    pub fn clear_indirect_branches(&self) {
+        unsafe { BNLowLevelILClearIndirectBranches(self.handle) };
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]