From 37f168e8f9e19c1dc133404b53a012fa4822bc9a Mon Sep 17 00:00:00 2001
From: Rubens Brandao <git@rubens.io>
Date: Fri, 31 May 2024 08:29:30 -0300
Subject: [PATCH] change `UserVariableValue` to the new standard

---
 rust/src/mlil/function.rs | 12 +++---
 rust/src/types.rs         | 87 +++++++++++++++++++++++++--------------
 2 files changed, 60 insertions(+), 39 deletions(-)

diff --git a/rust/src/mlil/function.rs b/rust/src/mlil/function.rs
index d9e133138..8c3ce3e23 100644
--- a/rust/src/mlil/function.rs
+++ b/rust/src/mlil/function.rs
@@ -11,7 +11,7 @@ use crate::function::{Function, Location};
 use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref, RefCountable};
 use crate::string::BnStrCompatible;
 use crate::types::{
-    Conf, PossibleValueSet, RegisterValue, SSAVariable, Type, UserVariableValues, Variable,
+    Conf, PossibleValueSet, RegisterValue, SSAVariable, Type, UserVariableValue, Variable,
 };
 
 use super::{MediumLevelILBlock, MediumLevelILInstruction, MediumLevelILLiftedInstruction};
@@ -242,20 +242,18 @@ impl MediumLevelILFunction {
 
     /// 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 {
+    pub fn user_var_values(&self) -> Array<UserVariableValue> {
         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),
-        }
+        unsafe { Array::new(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)?;
+        for user_value in self.user_var_values().iter() {
+            self.clear_user_var_value(&user_value.var, user_value.def_site.address)?;
         }
         Ok(())
     }
diff --git a/rust/src/types.rs b/rust/src/types.rs
index 59b24ef63..378c116ed 100644
--- a/rust/src/types.rs
+++ b/rust/src/types.rs
@@ -3355,55 +3355,78 @@ impl Drop for LookupTableEntryRaw {
 /////////////////////////
 // ArchAndAddr
 
+#[repr(C)]
 #[derive(Copy, Clone, Eq, Hash, PartialEq)]
 pub struct ArchAndAddr {
     pub arch: CoreArchitecture,
     pub address: u64,
 }
 
+impl ArchAndAddr {
+    pub(crate) unsafe fn from_raw(value: BNArchitectureAndAddress) -> Self {
+        mem::transmute(value)
+    }
+}
+
 /////////////////////////
-// UserVariableValues
+// UserVariableValue
+
+pub struct UserVariableValue {
+    pub var: Variable,
+    pub def_site: ArchAndAddr,
+    pub possible_value: PossibleValueSet,
+}
+
+impl UserVariableValue {
+    pub(crate) unsafe fn from_raw(value: BNUserVariableValue) -> Self {
+        let var = Variable::from_raw(value.var);
+        let def_site = ArchAndAddr::from_raw(value.defSite);
+        let value = PossibleValueSet::from_raw(value.value);
+        Self {
+            var,
+            def_site,
+            possible_value: value,
+        }
+    }
+}
 
-pub struct UserVariableValues {
-    pub(crate) vars: *const [BNUserVariableValue],
+impl CoreArrayProvider for UserVariableValue {
+    type Raw = BNUserVariableValue;
+    type Context = ();
+    type Wrapped<'a> = Self;
 }
 
-impl UserVariableValues {
-    pub fn into_hashmap(self) -> HashMap<Variable, HashMap<ArchAndAddr, PossibleValueSet>> {
+unsafe impl CoreArrayProviderInner for UserVariableValue {
+    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
+        unsafe { BNFreeUserVariableValues(raw) };
+    }
+
+    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
+        Self::from_raw(*raw)
+    }
+}
+
+impl Array<UserVariableValue> {
+    pub fn into_hashmap(&self) -> HashMap<Variable, HashMap<ArchAndAddr, PossibleValueSet>> {
         let mut result: HashMap<Variable, HashMap<ArchAndAddr, PossibleValueSet>> = HashMap::new();
-        for (var, def_site, possible_val) in self.all() {
+        for value in self.iter()
+        {
             result
-                .entry(var)
+                .entry(value.var)
                 .or_default()
-                .entry(def_site)
-                .or_insert(possible_val);
+                .entry(value.def_site)
+                .or_insert(value.possible_value);
         }
         result
     }
-    pub fn all(&self) -> impl Iterator<Item = (Variable, ArchAndAddr, PossibleValueSet)> {
-        unsafe { &*self.vars }.iter().map(|var_val| {
-            let var = unsafe { Variable::from_raw(var_val.var) };
-            let def_site = ArchAndAddr {
-                arch: unsafe { CoreArchitecture::from_raw(var_val.defSite.arch) },
-                address: var_val.defSite.address,
-            };
-            let possible_val = unsafe { PossibleValueSet::from_raw(var_val.value) };
-            (var, def_site, possible_val)
-        })
-    }
-    pub fn values_from_variable(
-        &self,
-        var: Variable,
-    ) -> impl Iterator<Item = (ArchAndAddr, PossibleValueSet)> {
-        self.all()
-            .filter(move |(t_var, _, _)| t_var == &var)
-            .map(|(_var, def_site, possible_val)| (def_site, possible_val))
-    }
-}
 
-impl Drop for UserVariableValues {
-    fn drop(&mut self) {
-        unsafe { BNFreeUserVariableValues(self.vars as *mut BNUserVariableValue) };
+    pub fn values_from_variable<'a>(
+        &'a self,
+        var: Variable,
+    ) -> impl Iterator<Item = (ArchAndAddr, PossibleValueSet)> + 'a {
+        self.iter()
+            .filter(move |value| value.var == var)
+            .map(|value| (value.def_site, value.possible_value))
     }
 }