add Function/MediumLevelILFunction user_var_values and related functions

rust/src/mlil/function.rs

@@ -1,19 +1,12 @@
 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;
-use binaryninjacore_sys::BNMediumLevelILFunction;
-use binaryninjacore_sys::BNMediumLevelILGetInstructionStart;
-use binaryninjacore_sys::BNNewMediumLevelILFunctionReference;
+use binaryninjacore_sys::*;
 
 use crate::basicblock::BasicBlock;
-use crate::function::Function;
-use crate::function::Location;
+use crate::function::{Function, Location};
 use crate::rc::{Array, Ref, RefCountable};
+use crate::string::BnStrCompatible;
+use crate::types::{Conf, PossibleValueSet, Type, UserVariableValues, Variable};
 
 use super::{MediumLevelILBlock, MediumLevelILInstruction, MediumLevelILLiftedInstruction};
 
@@ -105,6 +98,137 @@ impl MediumLevelILFunction {
 
         unsafe { Array::new(blocks, count, context) }
     }
+
+    pub fn get_var_definitions<'a>(&'a self, var: &Variable) -> VariableDefinitions<'a> {
+        let mut count = 0;
+        let raw_instrs =
+            unsafe { BNGetMediumLevelILVariableDefinitions(self.handle, &var.raw(), &mut count) };
+        assert!(!raw_instrs.is_null());
+        let instrs = unsafe { core::slice::from_raw_parts(raw_instrs, count) };
+        VariableDefinitions {
+            mlil: self,
+            ptr: raw_instrs,
+            instr_idxs: instrs.iter(),
+        }
+    }
+
+    pub fn create_user_var<'a, S: BnStrCompatible, C: Into<Conf<&'a Type>>>(
+        &self,
+        var: &Variable,
+        var_type: C,
+        name: S,
+        ignore_disjoint_uses: bool,
+    ) {
+        let var_type = var_type.into();
+        let raw_var_type: BNTypeWithConfidence = var_type.into();
+        let name = name.into_bytes_with_nul();
+        unsafe {
+            BNCreateUserVariable(
+                self.get_function().handle,
+                &var.raw(),
+                &raw_var_type as *const _ as *mut _,
+                name.as_ref().as_ptr() as *const _,
+                ignore_disjoint_uses,
+            )
+        }
+    }
+
+    pub fn delete_user_var(&self, var: &Variable) {
+        unsafe { BNDeleteUserVariable(self.get_function().handle, &var.raw()) }
+    }
+
+    pub fn is_var_user_defined(self, var: &Variable) -> bool {
+        unsafe { BNIsVariableUserDefined(self.get_function().handle, &var.raw()) }
+    }
+
+    /// Allows the user to specify a PossibleValueSet value for an MLIL
+    /// variable at its definition site.
+    ///
+    /// .. warning:: Setting the variable value, triggers a reanalysis of the
+    /// function and allows the dataflow to compute and propagate values which
+    /// depend on the current variable. This implies that branch conditions
+    /// whose values can be determined statically will be computed, leading to
+    /// potential branch elimination at the HLIL layer.
+    ///
+    /// * `var` - Variable for which the value is to be set
+    /// * `addr` - Address of the definition site of the variable
+    /// * `value` - Informed value of the variable
+    ///
+    /// # Example
+    /// ```no_run
+    /// # use binaryninja::function::Function;
+    /// # let fun: Function = todo!();
+    /// mlil_var = current_mlil[0].operands[0]
+    /// def_address = 0x40108d
+    /// var_value = PossibleValueSet.constant(5)
+    /// current_function.set_user_var_value(mlil_var, def_address, var_value)
+    /// ```
+    pub fn set_user_var_value(
+        &self,
+        var: &Variable,
+        addr: u64,
+        value: PossibleValueSet,
+    ) -> Result<(), ()> {
+        let Some(_def_site) = self
+            .get_var_definitions(var)
+            .find(|def| def.address == addr)
+        else {
+            // Error "No definition for Variable found at given address"
+            return Err(());
+        };
+        let function = self.get_function();
+        let def_site = BNArchitectureAndAddress {
+            arch: function.arch().0,
+            address: addr,
+        };
+        let value = value.into_raw();
+
+        unsafe { BNSetUserVariableValue(function.handle, &var.raw(), &def_site, value.as_ffi()) }
+        Ok(())
+    }
+
+    /// Clears a previously defined user variable value.
+    ///
+    /// * `var` - Variable for which the value was informed
+    /// * `def_addr` - Address of the definition site of the variable
+    pub fn clear_user_var_value(&self, var: &Variable, addr: u64) -> Result<(), ()> {
+        let Some(_var_def) = self
+            .get_var_definitions(var)
+            .find(|site| site.address == addr)
+        else {
+            //error "Could not get definition for Variable"
+            return Err(());
+        };
+
+        let function = self.get_function();
+        let def_site = BNArchitectureAndAddress {
+            arch: function.arch().0,
+            address: addr,
+        };
+
+        unsafe { BNClearUserVariableValue(function.handle, &var.raw(), &def_site) };
+        Ok(())
+    }
+
+    /// Returns a map of current defined user variable values.
+    /// Returns a Map of user current defined user variable values and their definition sites.
+    pub fn user_var_values(&self) -> UserVariableValues {
+        let mut count = 0;
+        let function = self.get_function();
+        let var_values = unsafe { BNGetAllUserVariableValues(function.handle, &mut count) };
+        assert!(!var_values.is_null());
+        UserVariableValues {
+            vars: core::ptr::slice_from_raw_parts(var_values, count),
+        }
+    }
+
+    /// Clear all user defined variable values.
+    pub fn clear_user_var_values(&self) -> Result<(), ()> {
+        for (var, arch_and_addr, _value) in self.user_var_values().all() {
+            self.clear_user_var_value(&var, arch_and_addr.address)?;
+        }
+        Ok(())
+    }
 }
 
 impl ToOwned for MediumLevelILFunction {
@@ -132,3 +256,37 @@ impl core::fmt::Debug for MediumLevelILFunction {
         write!(f, "<mlil func handle {:p}>", self.handle)
     }
 }
+
+#[derive(Clone, Debug)]
+pub struct VariableDefinitions<'a> {
+    mlil: &'a MediumLevelILFunction,
+    ptr: *mut usize,
+    instr_idxs: core::slice::Iter<'a, usize>,
+}
+
+impl Drop for VariableDefinitions<'_> {
+    fn drop(&mut self) {
+        unsafe { BNFreeILInstructionList(self.ptr) };
+    }
+}
+
+impl Iterator for VariableDefinitions<'_> {
+    type Item = MediumLevelILInstruction;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.instr_idxs
+            .next()
+            .map(|i| self.mlil.instruction_from_instruction_idx(*i))
+    }
+}
+
+impl DoubleEndedIterator for VariableDefinitions<'_> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.instr_idxs
+            .next_back()
+            .map(|i| self.mlil.instruction_from_instruction_idx(*i))
+    }
+}
+
+impl ExactSizeIterator for VariableDefinitions<'_> {}
+impl core::iter::FusedIterator for VariableDefinitions<'_> {}