Add After variants for looking up variables at instructions

Fixes #6397

2 files changed, 37 insertions, 0 deletions

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)