[Rust] Finish out the `PossibleValueSet` implementation to include range and lists of values

This requires us to allocate on the rust side of the API, this is quite easy to mess up but the usage of it is limited.

1 files changed, 72 insertions, 26 deletions

diff --git a/rust/src/variable.rs b/rust/src/variable.rs
index c0cab7d5..1921f950 100644
--- a/rust/src/variable.rs
+++ b/rust/src/variable.rs
@@ -12,7 +12,8 @@ use binaryninjacore_sys::{
     BNFreeVariableList, BNFreeVariableNameAndTypeList, BNFromVariableIdentifier,
     BNIndirectBranchInfo, BNLookupTableEntry, BNMergedVariable, BNPossibleValueSet,
     BNRegisterValue, BNRegisterValueType, BNStackVariableReference, BNToVariableIdentifier,
-    BNUserVariableValue, BNValueRange, BNVariable, BNVariableNameAndType, BNVariableSourceType,
+    BNTypeWithConfidence, BNUserVariableValue, BNValueRange, BNVariable, BNVariableNameAndType,
+    BNVariableSourceType,
 };
 use std::collections::HashSet;
 
@@ -238,7 +239,9 @@ impl UserVariableValue {
             var: value.variable.into(),
             defSite: value.def_site.into(),
             after: value.after,
-            value: PossibleValueSet::into_raw(value.value),
+            // TODO: This returns a rust allocated value, we should at some point provide allocators for the
+            // TODO: internal state of BNPossibleValueSet, so we can store rust created object in core objects.
+            value: PossibleValueSet::into_rust_raw(value.value),
         }
     }
 }
@@ -590,6 +593,28 @@ impl LookupTableEntry {
             to: value.toValue,
         }
     }
+
+    pub(crate) fn from_owned_raw(value: BNLookupTableEntry) -> Self {
+        let owned = Self::from_raw(&value);
+        Self::free_raw(value);
+        owned
+    }
+
+    pub(crate) fn into_raw(value: Self) -> BNLookupTableEntry {
+        let from_values: Box<[i64]> = value.from.into_iter().collect();
+        let from_values_len = from_values.len();
+        BNLookupTableEntry {
+            // Freed in [`Self::free_raw`]
+            fromValues: Box::leak(from_values).as_mut_ptr(),
+            fromCount: from_values_len,
+            toValue: value.to,
+        }
+    }
+
+    pub(crate) fn free_raw(value: BNLookupTableEntry) {
+        let raw_from = unsafe { std::slice::from_raw_parts_mut(value.fromValues, value.fromCount) };
+        let boxed_from = unsafe { Box::from_raw(raw_from) };
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -724,15 +749,14 @@ impl PossibleValueSet {
         }
     }
 
-    /// Take ownership over an "owned" core allocated value. Do not call this for a rust allocated value.
-    pub(crate) fn from_owned_raw(mut value: BNPossibleValueSet) -> Self {
+    /// Take ownership over an "owned" **core allocated** value. Do not call this for a rust allocated value.
+    pub(crate) fn from_owned_core_raw(mut value: BNPossibleValueSet) -> Self {
         let owned = Self::from_raw(&value);
-        // TODO: This entire function is a little wonky.
-        Self::free_raw(&mut value);
+        Self::free_core_raw(&mut value);
         owned
     }
 
-    pub(crate) fn into_raw(value: Self) -> BNPossibleValueSet {
+    pub(crate) fn into_rust_raw(value: Self) -> BNPossibleValueSet {
         let mut raw = BNPossibleValueSet {
             state: value.value_type(),
             ..Default::default()
@@ -758,31 +782,39 @@ impl PossibleValueSet {
             PossibleValueSet::ReturnAddressValue => {}
             PossibleValueSet::ImportedAddressValue => {}
             PossibleValueSet::SignedRangeValue { value, ranges } => {
+                let boxed_raw_ranges: Box<[BNValueRange]> =
+                    ranges.into_iter().map(BNValueRange::from).collect();
                 raw.value = value;
-                // TODO: raw.ranges
-                // TODO: requires core allocation and freeing.
-                // TODO: See `BNFreePossibleValueSet` for why this sucks.
+                raw.count = boxed_raw_ranges.len();
+                // NOTE: We are allocating this in rust, meaning core MUST NOT free this.
+                raw.ranges = Box::leak(boxed_raw_ranges).as_mut_ptr();
             }
             PossibleValueSet::UnsignedRangeValue { value, ranges } => {
+                let boxed_raw_ranges: Box<[BNValueRange]> =
+                    ranges.into_iter().map(BNValueRange::from).collect();
                 raw.value = value;
-                // TODO: raw.ranges
-                // TODO: requires core allocation and freeing.
-                // TODO: See `BNFreePossibleValueSet` for why this sucks.
+                raw.count = boxed_raw_ranges.len();
+                // NOTE: We are allocating this in rust, meaning core MUST NOT free this.
+                raw.ranges = Box::leak(boxed_raw_ranges).as_mut_ptr();
             }
             PossibleValueSet::LookupTableValue { table } => {
-                // TODO: raw.table
-                // TODO: requires core allocation and freeing.
-                // TODO: See `BNFreePossibleValueSet` for why this sucks.
+                let boxed_raw_entries: Box<[BNLookupTableEntry]> =
+                    table.into_iter().map(LookupTableEntry::into_raw).collect();
+                raw.count = boxed_raw_entries.len();
+                // NOTE: We are allocating this in rust, meaning core MUST NOT free this.
+                raw.table = Box::leak(boxed_raw_entries).as_mut_ptr();
             }
             PossibleValueSet::InSetOfValues { values } => {
-                // TODO: raw.valueSet
-                // TODO: requires core allocation and freeing.
-                // TODO: See `BNFreePossibleValueSet` for why this sucks.
+                let boxed_raw_values: Box<[i64]> = values.into_iter().collect();
+                raw.count = boxed_raw_values.len();
+                // NOTE: We are allocating this in rust, meaning core MUST NOT free this.
+                raw.valueSet = Box::leak(boxed_raw_values).as_mut_ptr();
             }
             PossibleValueSet::NotInSetOfValues { values } => {
-                // TODO: raw.valueSet
-                // TODO: requires core allocation and freeing.
-                // TODO: See `BNFreePossibleValueSet` for why this sucks.
+                let boxed_raw_values: Box<[i64]> = values.into_iter().collect();
+                raw.count = boxed_raw_values.len();
+                // NOTE: We are allocating this in rust, meaning core MUST NOT free this.
+                raw.valueSet = Box::leak(boxed_raw_values).as_mut_ptr();
             }
             PossibleValueSet::ConstantDataValue { value, size } => {
                 raw.value = value;
@@ -804,14 +836,28 @@ impl PossibleValueSet {
         raw
     }
 
-    /// Free a CORE ALLOCATED possible value set. Do not use this with [Self::into_raw] values.
-    pub(crate) fn free_raw(value: &mut BNPossibleValueSet) {
+    /// Free a CORE ALLOCATED possible value set. Do not use this with [Self::into_rust_raw] values.
+    pub(crate) fn free_core_raw(value: &mut BNPossibleValueSet) {
         unsafe { BNFreePossibleValueSet(value) }
     }
 
     /// Free a RUST ALLOCATED possible value set. Do not use this with CORE ALLOCATED values.
-    pub(crate) fn free_owned_raw(value: BNPossibleValueSet) {
-        // TODO: Once we fill out allocation of the possible value set then we should fill this out as well.
+    pub(crate) fn free_rust_raw(value: BNPossibleValueSet) {
+        // Free the range list
+        if !value.ranges.is_null() {
+            let raw_ranges = unsafe { std::slice::from_raw_parts_mut(value.ranges, value.count) };
+            let boxed_ranges = unsafe { Box::from_raw(raw_ranges) };
+        }
+
+        if !value.table.is_null() {
+            unsafe { LookupTableEntry::free_raw(*value.table) };
+        }
+
+        if !value.valueSet.is_null() {
+            let raw_value_set =
+                unsafe { std::slice::from_raw_parts_mut(value.valueSet, value.count) };
+            let boxed_value_set = unsafe { Box::from_raw(raw_value_set) };
+        }
     }
 
     pub fn value_type(&self) -> RegisterValueType {