[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.

4 files changed, 155 insertions, 45 deletions

diff --git a/rust/src/medium_level_il/function.rs b/rust/src/medium_level_il/function.rs
index 8751111b..4206f7dd 100644
--- a/rust/src/medium_level_il/function.rs
+++ b/rust/src/medium_level_il/function.rs
@@ -183,7 +183,7 @@ impl MediumLevelILFunction {
     /// 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
+    /// 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
@@ -212,23 +212,15 @@ impl MediumLevelILFunction {
         value: PossibleValueSet,
         after: bool,
     ) -> Result<(), ()> {
-        let Some(_def_site) = self
-            .var_definitions(var)
-            .iter()
-            .find(|def| def.address == addr)
-        else {
-            // Error "No definition for Variable found at given address"
-            return Err(());
-        };
         let function = self.function();
         let def_site = BNArchitectureAndAddress {
             arch: function.arch().handle,
             address: addr,
         };
         let raw_var = BNVariable::from(var);
-        let raw_value = PossibleValueSet::into_raw(value);
+        let raw_value = PossibleValueSet::into_rust_raw(value);
         unsafe { BNSetUserVariableValue(function.handle, &raw_var, &def_site, after, &raw_value) }
-        PossibleValueSet::free_owned_raw(raw_value);
+        PossibleValueSet::free_rust_raw(raw_value);
         Ok(())
     }
 
diff --git a/rust/src/medium_level_il/instruction.rs b/rust/src/medium_level_il/instruction.rs
index 3af9e424..ebe99275 100644
--- a/rust/src/medium_level_il/instruction.rs
+++ b/rust/src/medium_level_il/instruction.rs
@@ -1167,7 +1167,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     pub fn possible_ssa_variable_values(&self, ssa_var: &SSAVariable) -> PossibleValueSet {
@@ -1190,7 +1190,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     /// Return the ssa version of a [`Variable`] at the given instruction.
@@ -1390,7 +1390,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     pub fn possible_register_values_after(&self, reg_id: RegisterId) -> PossibleValueSet {
@@ -1411,7 +1411,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     pub fn flag_value(&self, flag_id: FlagId) -> RegisterValue {
@@ -1454,7 +1454,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     pub fn possible_flag_values_after_with_opts(
@@ -1471,7 +1471,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     pub fn stack_contents(&self, offset: i64, size: usize) -> RegisterValue {
@@ -1514,7 +1514,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     pub fn possible_stack_contents_after_with_opts(
@@ -1533,7 +1533,7 @@ impl MediumLevelILInstruction {
                 options.len(),
             )
         };
-        PossibleValueSet::from_owned_raw(value)
+        PossibleValueSet::from_owned_core_raw(value)
     }
 
     /// Gets the unique variable for a definition instruction. This unique variable can be passed
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 {
diff --git a/rust/tests/medium_level_il.rs b/rust/tests/medium_level_il.rs
index 7d4d8af1..de79a47d 100644
--- a/rust/tests/medium_level_il.rs
+++ b/rust/tests/medium_level_il.rs
@@ -4,6 +4,7 @@ use binaryninja::medium_level_il::{
     MediumLevelExpressionIndex, MediumLevelILInstructionKind, MediumLevelILLiftedInstructionKind,
     MediumLevelInstructionIndex,
 };
+use binaryninja::variable::{PossibleValueSet, ValueRange};
 use std::path::PathBuf;
 
 #[test]
@@ -117,3 +118,74 @@ fn test_mlil_basic_blocks() {
         }
     }
 }
+
+#[test]
+fn test_mlil_possible_values() {
+    let _session = Session::new().expect("Failed to initialize session");
+    let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
+    let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
+    let image_base = view.original_image_base();
+
+    let functions = view.functions_containing(image_base + 0x0002af9a);
+    assert_eq!(functions.len(), 1);
+    let function = functions.get(0);
+    let mlil_function = function.medium_level_il().unwrap();
+
+    // 0 @ 0002af40  (MLIL_SET_VAR.d edi_1 = (MLIL_VAR.d edi))
+    let instr_0 = mlil_function
+        .instruction_from_index(MediumLevelInstructionIndex(0))
+        .expect("Failed to get instruction");
+    let lifted_instr_0 = instr_0.lift();
+    match lifted_instr_0.kind {
+        MediumLevelILLiftedInstructionKind::SetVar(op) => match op.src.kind {
+            MediumLevelILLiftedInstructionKind::Var(var) => {
+                let var = var.src;
+                let pvs_value = PossibleValueSet::SignedRangeValue {
+                    value: 5,
+                    ranges: vec![ValueRange {
+                        start: -5,
+                        end: 0,
+                        step: 1,
+                    }],
+                };
+                mlil_function
+                    .set_user_var_value(&var, op.src.address, pvs_value.clone(), false)
+                    .expect("Failed to set user var value");
+                let var_values = mlil_function.user_var_values();
+                assert_eq!(var_values.len(), 1);
+                let var_value = var_values.get(0);
+                assert_eq!(var_value.value, pvs_value);
+                assert_eq!(var_value.def_site.addr, op.src.address);
+                assert_eq!(var_value.variable, var);
+                assert_eq!(var_value.after, false);
+            }
+            _ => panic!("Expected Var"),
+        },
+        _ => panic!("Expected SetVar"),
+    }
+
+    // 21 @ 0002af9a  (MLIL_RET return (MLIL_VAR.d eax_2))
+    let instr_21 = mlil_function
+        .instruction_from_index(MediumLevelInstructionIndex(21))
+        .expect("Failed to get instruction");
+    let lifted_instr_21 = instr_21.lift();
+    match lifted_instr_21.kind {
+        MediumLevelILLiftedInstructionKind::Ret(ret) => {
+            // This should be the eax variable, with the signed range.
+            let eax_var_lifted = ret.src.get(0).unwrap();
+            // TODO: I really dislike that we have lifted and unlifted versions of the same thing.
+            let eax_var = mlil_function
+                .instruction_from_expr_index(eax_var_lifted.expr_index)
+                .unwrap();
+            let eax_var_pvs = eax_var.possible_values();
+            match eax_var_pvs {
+                PossibleValueSet::SignedRangeValue { value, ranges } => {
+                    assert_eq!(value, -48);
+                    assert_eq!(ranges.len(), 2);
+                }
+                _ => panic!("Expected SignedRangeValue"),
+            }
+        }
+        _ => panic!("Expected Ret"),
+    }
+}