Resolves #2748; Variable.ssa_versions; Fixes some type annotations

5 files changed, 71 insertions, 30 deletions

diff --git a/python/callingconvention.py b/python/callingconvention.py
index 6ab587a0..5782c8e5 100644
--- a/python/callingconvention.py
+++ b/python/callingconvention.py
@@ -20,7 +20,7 @@
 
 import traceback
 import ctypes
-from typing import Optional
+from typing import Optional, Union
 
 # Binary Ninja components
 from . import _binaryninjacore as core
@@ -29,6 +29,11 @@ from . import variable
 from . import function
 from . import architecture
 
+FunctionOrILFunction = Union["binaryninja.function.Function",
+                             "binaryninja.lowlevelil.LowLevelILFunction",
+                             "binaryninja.mediumlevelil.MediumLevelILFunction",
+                             "binaryninja.highlevelil.HighLevelILFunction"]
+
 
 class CallingConvention:
 	name = None
@@ -439,7 +444,7 @@ class CallingConvention:
 			func_handle = func.handle
 		return variable.RegisterValue.from_BNRegisterValue(core.BNGetIncomingFlagValue(self.handle, reg_num, func_handle), self.arch)
 
-	def get_incoming_var_for_parameter_var(self, in_var:'variable.CoreVariable', func:'function.Function') -> 'variable.Variable':
+	def get_incoming_var_for_parameter_var(self, in_var:'variable.CoreVariable', func:FunctionOrILFunction) -> 'variable.Variable':
 		in_buf = in_var.to_BNVariable()
 		if func is None:
 			func_obj = None
@@ -448,7 +453,7 @@ class CallingConvention:
 		out_var = core.BNGetIncomingVariableForParameterVariable(self.handle, in_buf, func_obj)
 		return variable.Variable.from_BNVariable(func, out_var)
 
-	def get_parameter_var_for_incoming_var(self, in_var:'variable.CoreVariable', func:'function.Function') -> 'variable.Variable':
+	def get_parameter_var_for_incoming_var(self, in_var:'variable.CoreVariable', func:FunctionOrILFunction) -> 'variable.Variable':
 		in_buf = in_var.to_BNVariable()
 		if func is None:
 			func_obj = None
diff --git a/python/highlevelil.py b/python/highlevelil.py
index 128e50d8..a09da6f2 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -631,15 +631,15 @@ class HighLevelILInstruction(BaseILInstruction):
 
 	def get_var(self, operand_index:int) -> 'variable.Variable':
 		value = self.core_instr.operands[operand_index]
-		return variable.Variable.from_identifier(self.function.source_function, self.core_instr.operands[operand_index])
+		return variable.Variable.from_identifier(self.function, self.core_instr.operands[operand_index])
 
 	def get_var_ssa(self, operand_index1:int, operand_index2:int) -> 'mediumlevelil.SSAVariable':
-		var = variable.Variable.from_identifier(self.function.source_function, self.core_instr.operands[operand_index1])
+		var = variable.Variable.from_identifier(self.function, self.core_instr.operands[operand_index1])
 		version = self.core_instr.operands[operand_index2]
 		return mediumlevelil.SSAVariable(var, version)
 
 	def get_var_ssa_dest_and_src(self, operand_index1:int, operand_index2:int) -> 'mediumlevelil.SSAVariable':
-		var = variable.Variable.from_identifier(self.function.source_function, self.core_instr.operands[operand_index1])
+		var = variable.Variable.from_identifier(self.function, self.core_instr.operands[operand_index1])
 		dest_version = self.core_instr.operands[operand_index2]
 		return mediumlevelil.SSAVariable(var, dest_version)
 
@@ -677,7 +677,7 @@ class HighLevelILInstruction(BaseILInstruction):
 			for j in range(count.value // 2):
 				var_id = operand_list[j * 2]
 				var_version = operand_list[(j * 2) + 1]
-				value.append(mediumlevelil.SSAVariable(variable.Variable.from_identifier(self.function.source_function,
+				value.append(mediumlevelil.SSAVariable(variable.Variable.from_identifier(self.function,
 					var_id), var_version))
 			return value
 		finally:
@@ -2288,7 +2288,7 @@ class HighLevelILFunction:
 		var_data = ssa_var.var.to_BNVariable()
 		return core.BNIsHighLevelILSSAVarLive(self.handle, var_data, ssa_var.version)
 
-	def is_var_live_at(self, var: 'variable.Variable', instr: int) -> bool:
+	def is_var_live_at(self, var: 'variable.Variable', instr: InstructionIndex) -> bool:
 		"""
 		``is_var_live_at`` determines if ``var`` is live at a given point in the function
 		"""
@@ -2298,7 +2298,7 @@ class HighLevelILFunction:
 		var_data.storage = var.storage
 		return core.BNIsHighLevelILVarLiveAt(self.handle, var_data, instr)
 
-	def is_ssa_var_live_at(self, ssa_var: 'mediumlevelil.SSAVariable', instr: int) -> bool:
+	def is_ssa_var_live_at(self, ssa_var: 'mediumlevelil.SSAVariable', instr: InstructionIndex) -> bool:
 		"""
 		``is_ssa_var_live_at`` determines if ``ssa_var`` is live at a given point in the function; counts phi's as uses
 		"""
@@ -2378,14 +2378,14 @@ class HighLevelILFunction:
 		if self.source_function is None:
 			return []
 
-		if self.il_form == FunctionGraphType.HighLevelILFunctionGraph or self.il_form == FunctionGraphType.HighLevelILSSAFormFunctionGraph:
+		if self.il_form in [FunctionGraphType.HighLevelILFunctionGraph, FunctionGraphType.HighLevelILSSAFormFunctionGraph]:
 			count = ctypes.c_ulonglong()
 			core_variables = core.BNGetHighLevelILVariables(self.handle, count)
 			assert core_variables is not None, "core.BNGetHighLevelILVariables returned None"
 			try:
 				result = []
 				for var_i in range(count.value):
-					result.append(variable.Variable(self.source_function, core_variables[var_i].type, core_variables[var_i].index, core_variables[var_i].storage))
+					result.append(variable.Variable(self, core_variables[var_i].type, core_variables[var_i].index, core_variables[var_i].storage))
 				return result
 			finally:
 				core.BNFreeVariableList(core_variables)
@@ -2409,7 +2409,7 @@ class HighLevelILFunction:
 					assert versions is not None, "core.BNGetHighLevelILVariableSSAVersions returned None"
 					try:
 						for version_i in range(version_count.value):
-							result.append(mediumlevelil.SSAVariable(variable.Variable(self.source_function, core_variables[var_i].type, core_variables[var_i].index, core_variables[var_i].storage), versions[version_i]))
+							result.append(mediumlevelil.SSAVariable(variable.Variable(self, core_variables[var_i].type, core_variables[var_i].index, core_variables[var_i].storage), versions[version_i]))
 					finally:
 						core.BNFreeILInstructionList(versions)
 				return result
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 42c2878c..50e058e9 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -3008,7 +3008,8 @@ class LowLevelILFunction:
 		return []
 
 	@property
-	def ssa_vars(self) -> List[Union['mediumlevelil.SSAVariable', SSARegisterStack, SSAFlag]]:
+	def ssa_vars(self) -> List[Union[SSARegister, SSARegisterStack, SSAFlag]]:
+		# TODO : Add ssa_varsions to these too
 		"""This is the union `LowLevelILFunction.ssa_registers`, `LowLevelILFunction.ssa_register_stacks`, and `LowLevelILFunction.ssa_flags`"""
 		if self.il_form == FunctionGraphType.LowLevelILSSAFormFunctionGraph:
 			return self.ssa_registers + self.ssa_register_stacks + self.ssa_flags  # type: ignore
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index d0a5bc8e..8f079bfe 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -524,16 +524,16 @@ class MediumLevelILInstruction(BaseILInstruction):
 	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.source_function, result)
+		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.source_function, result)
+		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.source_function, result)
+		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)
@@ -698,15 +698,15 @@ class MediumLevelILInstruction(BaseILInstruction):
 
 	def get_var(self, operand_index:int) -> variable.Variable:
 		value = self.instr.operands[operand_index]
-		return variable.Variable.from_identifier(self.function.source_function, value)
+		return variable.Variable.from_identifier(self.function, value)
 
 	def get_var_ssa(self, operand_index1:int, operand_index2:int) -> SSAVariable:
-		var = variable.Variable.from_identifier(self.function.source_function, self.instr.operands[operand_index1])
+		var = variable.Variable.from_identifier(self.function, self.instr.operands[operand_index1])
 		version = self.instr.operands[operand_index2]
 		return SSAVariable(var, version)
 
 	def get_var_ssa_dest_and_src(self, operand_index1:int, operand_index2:int) -> SSAVariable:
-		var = variable.Variable.from_identifier(self.function.source_function, self.instr.operands[operand_index1])
+		var = variable.Variable.from_identifier(self.function, self.instr.operands[operand_index1])
 		dest_version = self.instr.operands[operand_index2]
 		return SSAVariable(var, dest_version)
 
@@ -729,7 +729,7 @@ class MediumLevelILInstruction(BaseILInstruction):
 		value:List[variable.Variable] = []
 		try:
 			for j in range(count.value):
-				value.append(variable.Variable.from_identifier(self.function.source_function, operand_list[j]))
+				value.append(variable.Variable.from_identifier(self.function, operand_list[j]))
 			return value
 		finally:
 			core.BNMediumLevelILFreeOperandList(operand_list)
@@ -743,7 +743,7 @@ class MediumLevelILInstruction(BaseILInstruction):
 			for j in range(count.value // 2):
 				var_id = operand_list[j * 2]
 				var_version = operand_list[(j * 2) + 1]
-				value.append(SSAVariable(variable.Variable.from_identifier(self.function.source_function,
+				value.append(SSAVariable(variable.Variable.from_identifier(self.function,
 					var_id), var_version))
 			return value
 		finally:
@@ -3030,7 +3030,7 @@ class MediumLevelILFunction:
 			result = []
 			try:
 				for var_i in range(count.value):
-					result.append(variable.Variable(self.source_function, core_variables[var_i].type, core_variables[var_i].index, core_variables[var_i].storage))
+					result.append(variable.Variable(self, core_variables[var_i].type, core_variables[var_i].index, core_variables[var_i].storage))
 				return result
 			finally:
 				core.BNFreeVariableList(core_variables)
@@ -3054,7 +3054,7 @@ class MediumLevelILFunction:
 					assert versions is not None, "core.BNGetMediumLevelILVariableSSAVersions returned None"
 					try:
 						for version_i in range(version_count.value):
-							result.append(SSAVariable(variable.Variable(self.source_function,
+							result.append(SSAVariable(variable.Variable(self,
 								core_variables[var_i].type, core_variables[var_i].index,
 								core_variables[var_i].storage), versions[version_i]))
 					finally:
diff --git a/python/variable.py b/python/variable.py
index d345e4b0..3ce7e437 100644
--- a/python/variable.py
+++ b/python/variable.py
@@ -26,7 +26,13 @@ from dataclasses import dataclass, field
 import binaryninja
 from . import _binaryninjacore as core
 from . import decorators
-from .enums import RegisterValueType, VariableSourceType, DeadStoreElimination
+from .enums import RegisterValueType, VariableSourceType, DeadStoreElimination, FunctionGraphType
+
+FunctionOrILFunction = Union["binaryninja.function.Function",
+                             "binaryninja.lowlevelil.LowLevelILFunction",
+                             "binaryninja.mediumlevelil.MediumLevelILFunction",
+                             "binaryninja.highlevelil.HighLevelILFunction"]
+
 
 @dataclass(frozen=True)
 class LookupTableEntry:
@@ -624,24 +630,29 @@ class VariableNameAndType(CoreVariable):
 
 
 class Variable(CoreVariable):
-	def __init__(self, func:'binaryninja.function.Function', source_type:VariableSourceType, index:int, storage:int):
+	def __init__(self, func:FunctionOrILFunction, source_type:VariableSourceType, index:int, storage:int):
 		super(Variable, self).__init__(source_type, index, storage)
-		self._function = func
+		if isinstance(func, binaryninja.function.Function):
+			self._function = func
+			self._il_function = None
+		else:
+			self._function = func.source_function
+			self._il_function = func
 
 	@classmethod
-	def from_variable_name_and_type(cls, func:'binaryninja.function.Function', var:VariableNameAndType):
+	def from_variable_name_and_type(cls, func:FunctionOrILFunction, var:VariableNameAndType):
 		return cls(func, VariableSourceType(var.type), var.index, var.storage)
 
 	@classmethod
-	def from_core_variable(cls, func:'binaryninja.function.Function', var:CoreVariable):
+	def from_core_variable(cls, func:FunctionOrILFunction, var:CoreVariable):
 		return cls(func, var.source_type, var.index, var.storage)
 
 	@classmethod
-	def from_BNVariable(cls, func:'binaryninja.function.Function', var:core.BNVariable):
+	def from_BNVariable(cls, func:FunctionOrILFunction, var:core.BNVariable):
 		return cls(func, var.type, var.index, var.storage)
 
 	@classmethod
-	def from_identifier(cls, func:'binaryninja.function.Function', identifier:int):
+	def from_identifier(cls, func:FunctionOrILFunction, identifier:int):
 		var = core.BNFromVariableIdentifier(identifier)
 		return cls(func, VariableSourceType(var.type), var.index, var.storage)
 
@@ -716,6 +727,30 @@ class Variable(CoreVariable):
 		self._function.create_user_var(self, new_type, self.name)
 
 	@property
+	def ssa_versions(self) -> Generator[int, None, None]:
+		"""Returns the SSA versions associated with this variable.  Doesn't return anything for aliased variables."""
+		if self._il_function is None:
+			raise NotImplementedError("No IL function associated with variable")
+
+		version_count = ctypes.c_ulonglong()
+		versions = None
+		if self._il_function.il_form in [FunctionGraphType.MediumLevelILFunctionGraph, FunctionGraphType.MediumLevelILSSAFormFunctionGraph]:
+			versions = core.BNGetMediumLevelILVariableSSAVersions(self._il_function.handle, self.to_BNVariable(), version_count)
+		elif self._il_function.il_form in [FunctionGraphType.HighLevelILFunctionGraph, FunctionGraphType.HighLevelILSSAFormFunctionGraph]:
+			versions = core.BNGetHighLevelILVariableSSAVersions(self._il_function.handle, self.to_BNVariable(), version_count)
+		else:
+			raise NotImplementedError("Unsupported IL form")
+
+		if versions is None:
+			raise NotImplementedError("No SSA versions; is this an aliased variable?")
+
+		try:
+			for version_i in range(version_count.value):
+				yield versions[version_i]
+		finally:
+			core.BNFreeILInstructionList(versions)
+
+	@property
 	def dead_store_elimination(self):
 		return DeadStoreElimination(core.BNGetFunctionVariableDeadStoreElimination(self._function.handle, self.to_BNVariable()))