Add After variants for looking up variables at instructions

Fixes #6397

2 files changed, 85 insertions, 0 deletions

diff --git a/rust/src/function.rs b/rust/src/function.rs
index f122aeff..b438520a 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -1944,6 +1944,35 @@ impl Function {
         Some((var, name, var_type))
     }
 
+    pub fn stack_var_at_frame_offset_after_instruction(
+        &self,
+        addr: u64,
+        offset: i64,
+        arch: Option<CoreArchitecture>,
+    ) -> Option<(Variable, BnString, Conf<Ref<Type>>)> {
+        let arch = arch.unwrap_or_else(|| self.arch());
+        let mut found_value = BNVariableNameAndType::default();
+        let found = unsafe {
+            BNGetStackVariableAtFrameOffsetAfterInstruction(
+                self.handle,
+                arch.handle,
+                addr,
+                offset,
+                &mut found_value,
+            )
+        };
+        if !found {
+            return None;
+        }
+        let var = Variable::from(found_value.var);
+        let name = unsafe { BnString::from_raw(found_value.name) };
+        let var_type = Conf::new(
+            unsafe { Type::ref_from_raw(found_value.type_) },
+            found_value.typeConfidence,
+        );
+        Some((var, name, var_type))
+    }
+
     pub fn stack_variables_referenced_by(
         &self,
         addr: u64,
diff --git a/rust/src/medium_level_il/instruction.rs b/rust/src/medium_level_il/instruction.rs
index bd07da94..0097e83b 100644
--- a/rust/src/medium_level_il/instruction.rs
+++ b/rust/src/medium_level_il/instruction.rs
@@ -1029,6 +1029,19 @@ impl MediumLevelILInstruction {
         SSAVariable::new(var, version)
     }
 
+    /// Return the ssa version of a [`Variable`] after the given instruction.
+    pub fn ssa_variable_version_after(&self, var: Variable) -> SSAVariable {
+        let raw_var = BNVariable::from(var);
+        let version = unsafe {
+            BNGetMediumLevelILSSAVarVersionAfterILInstruction(
+                self.function.handle,
+                &raw_var,
+                self.expr_index.0,
+            )
+        };
+        SSAVariable::new(var, version)
+    }
+
     /// Set of branching instructions that must take the true or false path to reach this instruction
     pub fn branch_dependencies(&self) -> Array<BranchDependence> {
         let mut count = 0;
@@ -1067,6 +1080,16 @@ impl MediumLevelILInstruction {
         }
     }
 
+    /// Version of active memory contents in SSA form for this instruction
+    pub fn ssa_memory_version_after(&self) -> usize {
+        unsafe {
+            BNGetMediumLevelILSSAMemoryVersionAfterILInstruction(
+                self.function.handle,
+                self.expr_index.0,
+            )
+        }
+    }
+
     /// Type of expression
     pub fn expr_type(&self) -> Option<Conf<Ref<Type>>> {
         let result = unsafe { BNGetMediumLevelILExprType(self.function.handle, self.expr_index.0) };
@@ -1095,6 +1118,17 @@ impl MediumLevelILInstruction {
         Variable::from(result)
     }
 
+    pub fn variable_for_register_after(&self, reg_id: RegisterId) -> Variable {
+        let result = unsafe {
+            BNGetMediumLevelILVariableForRegisterAfterInstruction(
+                self.function.handle,
+                reg_id.0,
+                self.expr_index.0,
+            )
+        };
+        Variable::from(result)
+    }
+
     pub fn variable_for_flag(&self, flag_id: FlagId) -> Variable {
         let result = unsafe {
             BNGetMediumLevelILVariableForFlagAtInstruction(
@@ -1106,6 +1140,17 @@ impl MediumLevelILInstruction {
         Variable::from(result)
     }
 
+    pub fn variable_for_flag_after(&self, flag_id: FlagId) -> Variable {
+        let result = unsafe {
+            BNGetMediumLevelILVariableForFlagAfterInstruction(
+                self.function.handle,
+                flag_id.0,
+                self.expr_index.0,
+            )
+        };
+        Variable::from(result)
+    }
+
     pub fn variable_for_stack_location(&self, offset: i64) -> Variable {
         let result = unsafe {
             BNGetMediumLevelILVariableForStackLocationAtInstruction(
@@ -1117,6 +1162,17 @@ impl MediumLevelILInstruction {
         Variable::from(result)
     }
 
+    pub fn variable_for_stack_location_after(&self, offset: i64) -> Variable {
+        let result = unsafe {
+            BNGetMediumLevelILVariableForStackLocationAfterInstruction(
+                self.function.handle,
+                offset,
+                self.expr_index.0,
+            )
+        };
+        Variable::from(result)
+    }
+
     pub fn register_value(&self, reg_id: RegisterId) -> RegisterValue {
         unsafe {
             BNGetMediumLevelILRegisterValueAtInstruction(