Add After variants for looking up variables at instructions

Fixes #6397

8 files changed, 188 insertions, 1 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 8805efd7..45f797ca 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -10956,6 +10956,7 @@ namespace BinaryNinja {
 		void DeleteAutoStackVariable(int64_t offset);
 		void DeleteUserStackVariable(int64_t offset);
 		bool GetStackVariableAtFrameOffset(Architecture* arch, uint64_t addr, int64_t offset, VariableNameAndType& var);
+		bool GetStackVariableAtFrameOffsetAfterInstruction(Architecture* arch, uint64_t addr, int64_t offset, VariableNameAndType& var);
 
 		/*! List of Function Variables
 
@@ -13407,10 +13408,15 @@ namespace BinaryNinja {
 		    const std::set<BNDataFlowQueryOption>& options = std::set<BNDataFlowQueryOption>());
 
 		size_t GetSSAVarVersionAtInstruction(const Variable& var, size_t instr) const;
+		size_t GetSSAVarVersionAfterInstruction(const Variable& var, size_t instr) const;
 		size_t GetSSAMemoryVersionAtInstruction(size_t instr) const;
+		size_t GetSSAMemoryVersionAfterInstruction(size_t instr) const;
 		Variable GetVariableForRegisterAtInstruction(uint32_t reg, size_t instr) const;
+		Variable GetVariableForRegisterAfterInstruction(uint32_t reg, size_t instr) const;
 		Variable GetVariableForFlagAtInstruction(uint32_t flag, size_t instr) const;
+		Variable GetVariableForFlagAfterInstruction(uint32_t flag, size_t instr) const;
 		Variable GetVariableForStackLocationAtInstruction(int64_t offset, size_t instr) const;
+		Variable GetVariableForStackLocationAfterInstruction(int64_t offset, size_t instr) const;
 
 		RegisterValue GetRegisterValueAtInstruction(uint32_t reg, size_t instr);
 		RegisterValue GetRegisterValueAfterInstruction(uint32_t reg, size_t instr);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index ccc40bdf..50610112 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,7 +37,7 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 92
+#define BN_CURRENT_CORE_ABI_VERSION 93
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
@@ -4922,6 +4922,8 @@ extern "C"
 	BINARYNINJACOREAPI void BNDeleteUserStackVariable(BNFunction* func, int64_t offset);
 	BINARYNINJACOREAPI bool BNGetStackVariableAtFrameOffset(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, int64_t offset, BNVariableNameAndType* var);
+	BINARYNINJACOREAPI bool BNGetStackVariableAtFrameOffsetAfterInstruction(
+	    BNFunction* func, BNArchitecture* arch, uint64_t addr, int64_t offset, BNVariableNameAndType* var);
 	BINARYNINJACOREAPI void BNFreeVariableNameAndType(BNVariableNameAndType* var);
 
 	BINARYNINJACOREAPI BNVariableNameAndType* BNGetFunctionVariables(BNFunction* func, size_t* count);
@@ -5962,14 +5964,24 @@ extern "C"
 
 	BINARYNINJACOREAPI size_t BNGetMediumLevelILSSAVarVersionAtILInstruction(
 	    BNMediumLevelILFunction* func, const BNVariable* var, size_t instr);
+	BINARYNINJACOREAPI size_t BNGetMediumLevelILSSAVarVersionAfterILInstruction(
+	    BNMediumLevelILFunction* func, const BNVariable* var, size_t instr);
 	BINARYNINJACOREAPI size_t BNGetMediumLevelILSSAMemoryVersionAtILInstruction(
 	    BNMediumLevelILFunction* func, size_t instr);
+	BINARYNINJACOREAPI size_t BNGetMediumLevelILSSAMemoryVersionAfterILInstruction(
+	    BNMediumLevelILFunction* func, size_t instr);
 	BINARYNINJACOREAPI BNVariable BNGetMediumLevelILVariableForRegisterAtInstruction(
 	    BNMediumLevelILFunction* func, uint32_t reg, size_t instr);
+	BINARYNINJACOREAPI BNVariable BNGetMediumLevelILVariableForRegisterAfterInstruction(
+	    BNMediumLevelILFunction* func, uint32_t reg, size_t instr);
 	BINARYNINJACOREAPI BNVariable BNGetMediumLevelILVariableForFlagAtInstruction(
 	    BNMediumLevelILFunction* func, uint32_t flag, size_t instr);
+	BINARYNINJACOREAPI BNVariable BNGetMediumLevelILVariableForFlagAfterInstruction(
+	    BNMediumLevelILFunction* func, uint32_t flag, size_t instr);
 	BINARYNINJACOREAPI BNVariable BNGetMediumLevelILVariableForStackLocationAtInstruction(
 	    BNMediumLevelILFunction* func, int64_t offset, size_t instr);
+	BINARYNINJACOREAPI BNVariable BNGetMediumLevelILVariableForStackLocationAfterInstruction(
+	    BNMediumLevelILFunction* func, int64_t offset, size_t instr);
 
 	BINARYNINJACOREAPI BNRegisterValue BNGetMediumLevelILRegisterValueAtInstruction(
 	    BNMediumLevelILFunction* func, uint32_t reg, size_t instr);
diff --git a/function.cpp b/function.cpp
index 6c74dfc9..811e64c3 100644
--- a/function.cpp
+++ b/function.cpp
@@ -1358,6 +1358,23 @@ bool Function::GetStackVariableAtFrameOffset(
 }
 
 
+bool Function::GetStackVariableAtFrameOffsetAfterInstruction(
+    Architecture* arch, uint64_t addr, int64_t offset, VariableNameAndType& result)
+{
+	BNVariableNameAndType var;
+	if (!BNGetStackVariableAtFrameOffsetAfterInstruction(m_object, arch->GetObject(), addr, offset, &var))
+		return false;
+
+	result.type = Confidence<Ref<Type>>(new Type(BNNewTypeReference(var.type)), var.typeConfidence);
+	result.name = var.name;
+	result.var = var.var;
+	result.autoDefined = var.autoDefined;
+
+	BNFreeVariableNameAndType(&var);
+	return true;
+}
+
+
 map<Variable, VariableNameAndType> Function::GetVariables()
 {
 	size_t count;
diff --git a/mediumlevelil.cpp b/mediumlevelil.cpp
index a0a0b9dd..13a1a959 100644
--- a/mediumlevelil.cpp
+++ b/mediumlevelil.cpp
@@ -617,30 +617,60 @@ size_t MediumLevelILFunction::GetSSAVarVersionAtInstruction(const Variable& var,
 }
 
 
+size_t MediumLevelILFunction::GetSSAVarVersionAfterInstruction(const Variable& var, size_t instr) const
+{
+	return BNGetMediumLevelILSSAVarVersionAfterILInstruction(m_object, &var, instr);
+}
+
+
 size_t MediumLevelILFunction::GetSSAMemoryVersionAtInstruction(size_t instr) const
 {
 	return BNGetMediumLevelILSSAMemoryVersionAtILInstruction(m_object, instr);
 }
 
 
+size_t MediumLevelILFunction::GetSSAMemoryVersionAfterInstruction(size_t instr) const
+{
+	return BNGetMediumLevelILSSAMemoryVersionAfterILInstruction(m_object, instr);
+}
+
+
 Variable MediumLevelILFunction::GetVariableForRegisterAtInstruction(uint32_t reg, size_t instr) const
 {
 	return BNGetMediumLevelILVariableForRegisterAtInstruction(m_object, reg, instr);
 }
 
 
+Variable MediumLevelILFunction::GetVariableForRegisterAfterInstruction(uint32_t reg, size_t instr) const
+{
+	return BNGetMediumLevelILVariableForRegisterAfterInstruction(m_object, reg, instr);
+}
+
+
 Variable MediumLevelILFunction::GetVariableForFlagAtInstruction(uint32_t flag, size_t instr) const
 {
 	return BNGetMediumLevelILVariableForFlagAtInstruction(m_object, flag, instr);
 }
 
 
+Variable MediumLevelILFunction::GetVariableForFlagAfterInstruction(uint32_t flag, size_t instr) const
+{
+	return BNGetMediumLevelILVariableForFlagAfterInstruction(m_object, flag, instr);
+}
+
+
 Variable MediumLevelILFunction::GetVariableForStackLocationAtInstruction(int64_t offset, size_t instr) const
 {
 	return BNGetMediumLevelILVariableForStackLocationAtInstruction(m_object, offset, instr);
 }
 
 
+Variable MediumLevelILFunction::GetVariableForStackLocationAfterInstruction(int64_t offset, size_t instr) const
+{
+	return BNGetMediumLevelILVariableForStackLocationAfterInstruction(m_object, offset, instr);
+}
+
+
 RegisterValue MediumLevelILFunction::GetRegisterValueAtInstruction(uint32_t reg, size_t instr)
 {
 	BNRegisterValue value = BNGetMediumLevelILRegisterValueAtInstruction(m_object, reg, instr);
diff --git a/python/function.py b/python/function.py
index 14cc5e43..d4c7c92c 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2638,6 +2638,18 @@ class Function:
 		core.BNFreeVariableNameAndType(found_var)
 		return result
 
+	def get_stack_var_at_frame_offset_after_instruction(
+	    self, offset: int, addr: int, arch: Optional['architecture.Architecture'] = None
+	) -> Optional['variable.Variable']:
+		if arch is None:
+			arch = self.arch
+		found_var = core.BNVariableNameAndType()
+		if not core.BNGetStackVariableAtFrameOffsetAfterInstruction(self.handle, arch.handle, addr, offset, found_var):
+			return None
+		result = variable.Variable.from_BNVariable(self, found_var.var)
+		core.BNFreeVariableNameAndType(found_var)
+		return result
+
 	def get_type_tokens(self, settings: Optional['DisassemblySettings'] = None) -> List['DisassemblyTextLine']:
 		_settings = None
 		if settings is not None:
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 267076db..976c1a85 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -699,6 +699,11 @@ class MediumLevelILInstruction(BaseILInstruction):
 		return core.BNGetMediumLevelILSSAMemoryVersionAtILInstruction(self.function.handle, self.instr_index)
 
 	@property
+	def ssa_memory_version_after(self) -> int:
+		"""Version of active memory contents in SSA form after this instruction"""
+		return core.BNGetMediumLevelILSSAMemoryVersionAfterILInstruction(self.function.handle, self.instr_index)
+
+	@property
 	def prefix_operands(self) -> List[MediumLevelILOperandType]:
 		"""All operands in the expression tree in prefix order"""
 		result: List[MediumLevelILOperandType] = [MediumLevelILOperationAndSize(self.operation, self.size)]
@@ -803,22 +808,42 @@ class MediumLevelILInstruction(BaseILInstruction):
 		var_data = var.to_BNVariable()
 		return core.BNGetMediumLevelILSSAVarVersionAtILInstruction(self.function.handle, var_data, self.instr_index)
 
+	def get_ssa_var_version_after(self, var: variable.Variable) -> int:
+		var_data = var.to_BNVariable()
+		return core.BNGetMediumLevelILSSAVarVersionAfterILInstruction(self.function.handle, var_data, self.instr_index)
+
 	def get_var_for_reg(self, reg: 'architecture.RegisterType') -> variable.Variable:
 		reg = self.function.arch.get_reg_index(reg)
 		result = core.BNGetMediumLevelILVariableForRegisterAtInstruction(self.function.handle, reg, self.instr_index)
 		return variable.Variable.from_BNVariable(self.function, result)
 
+	def get_var_for_reg_after(self, reg: 'architecture.RegisterType') -> variable.Variable:
+		reg = self.function.arch.get_reg_index(reg)
+		result = core.BNGetMediumLevelILVariableForRegisterAfterInstruction(self.function.handle, reg, self.instr_index)
+		return variable.Variable.from_BNVariable(self.function, result)
+
 	def get_var_for_flag(self, flag: 'architecture.FlagType') -> variable.Variable:
 		flag = self.function.arch.get_flag_index(flag)
 		result = core.BNGetMediumLevelILVariableForFlagAtInstruction(self.function.handle, flag, self.instr_index)
 		return variable.Variable.from_BNVariable(self.function, result)
 
+	def get_var_for_flag_after(self, flag: 'architecture.FlagType') -> variable.Variable:
+		flag = self.function.arch.get_flag_index(flag)
+		result = core.BNGetMediumLevelILVariableForFlagAfterInstruction(self.function.handle, flag, self.instr_index)
+		return variable.Variable.from_BNVariable(self.function, result)
+
 	def get_var_for_stack_location(self, offset: int) -> variable.Variable:
 		result = core.BNGetMediumLevelILVariableForStackLocationAtInstruction(
 		    self.function.handle, offset, self.instr_index
 		)
 		return variable.Variable.from_BNVariable(self.function, result)
 
+	def get_var_for_stack_location_after(self, offset: int) -> variable.Variable:
+		result = core.BNGetMediumLevelILVariableForStackLocationAfterInstruction(
+		    self.function.handle, offset, self.instr_index
+		)
+		return variable.Variable.from_BNVariable(self.function, result)
+
 	def get_reg_value(self, reg: 'architecture.RegisterType') -> 'variable.RegisterValue':
 		reg = self.function.arch.get_reg_index(reg)
 		value = core.BNGetMediumLevelILRegisterValueAtInstruction(self.function.handle, reg, self.instr_index)
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(