Use new variable system in functions

2 files changed, 44 insertions, 25 deletions

diff --git a/python/function.py b/python/function.py
index 51d61792..74f4b0f8 100644
--- a/python/function.py
+++ b/python/function.py
@@ -26,7 +26,7 @@ import ctypes
 import _binaryninjacore as core
 from enums import (FunctionGraphType, BranchType, SymbolType, InstructionTextTokenType,
 	HighlightStandardColor, HighlightColorStyle, RegisterValueType, ImplicitRegisterExtend,
-	DisassemblyOption, IntegerDisplayType, InstructionTextTokenContext)
+	DisassemblyOption, IntegerDisplayType, InstructionTextTokenContext, VariableSourceType)
 import architecture
 import highlight
 import associateddatastore
@@ -146,14 +146,20 @@ class PossibleValueSet(object):
 		return "<undetermined>"
 
 
-class StackVariable(object):
-	def __init__(self, ofs, name, t):
-		self.offset = ofs
+class VariableNameAndType(object):
+	def __init__(self, var, name, t):
+		self.var = var
 		self.name = name
 		self.type = t
 
 	def __repr__(self):
-		return "<var@%x: %s %s>" % (self.offset, self.type, self.name)
+		if self.var.type == VariableSourceType.StackVariableSourceType:
+			return "<stack var %x: %s %s>" % (self.var.identifier, self.type, self.name)
+		elif self.var.type == VariableSourceType.RegisterVariableSourceType:
+			return "<reg var %s: %s %s>" % (self.var.function.arch.get_reg_name(self.var.identifier), self.type, self.name)
+		elif self.var.type == VariableSourceType.FlagVariableSourceType:
+			return "<flag var %s: %s %s>" % (self.var.function.arch.get_flag_name(self.var.identifier), self.type, self.name)
+		return "<var %s: %s %s>" % (self.var, self.type, self.name)
 
 	def __str__(self):
 		return self.name
@@ -179,7 +185,7 @@ class StackVariableReference(object):
 		return "<operand %d ref to %s>" % (self.source_operand, self.name)
 
 
-class ILVariable(object):
+class Variable(object):
 	def __init__(self, func, var_type, index, identifier):
 		self.function = func
 		self.type = var_type
@@ -363,14 +369,27 @@ class Function(object):
 
 	@property
 	def stack_layout(self):
-		"""List of function stack (read-only)"""
+		"""List of function stack variables (read-only)"""
 		count = ctypes.c_ulonglong()
 		v = core.BNGetStackLayout(self.handle, count)
 		result = []
 		for i in xrange(0, count.value):
-			result.append(StackVariable(v[i].offset, v[i].name, types.Type(handle = core.BNNewTypeReference(v[i].type))))
-		result.sort(key = lambda x: x.offset)
-		core.BNFreeStackLayout(v, count.value)
+			var = Variable(self, v[i].var.type, v[i].var.index, v[i].var.identifier)
+			result.append(VariableNameAndType(var, v[i].name, types.Type(handle = core.BNNewTypeReference(v[i].type))))
+		result.sort(key = lambda x: x.var.identifier)
+		core.BNFreeVariableList(v, count.value)
+		return result
+
+	@property
+	def vars(self):
+		"""List of function variables (read-only)"""
+		count = ctypes.c_ulonglong()
+		v = core.BNGetFunctionVariables(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			var = Variable(self, v[i].var.type, v[i].var.index, v[i].var.identifier)
+			result.append(VariableNameAndType(var, v[i].name, types.Type(handle = core.BNNewTypeReference(v[i].type))))
+		core.BNFreeVariableList(v, count.value)
 		return result
 
 	@property
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index dff79fbd..f827819f 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -22,7 +22,7 @@ import ctypes
 
 # Binary Ninja components
 import _binaryninjacore as core
-from .enums import MediumLevelILOperation, InstructionTextTokenType, ILVariableSourceType, ILBranchDependence
+from .enums import MediumLevelILOperation, InstructionTextTokenType, VariableSourceType, ILBranchDependence
 import function
 import basicblock
 import lowlevelil
@@ -157,11 +157,11 @@ class MediumLevelILInstruction(object):
 			elif operand_type == "expr":
 				value = MediumLevelILInstruction(func, instr.operands[i])
 			elif operand_type == "var":
-				var_type = ILVariableSourceType(instr.operands[i] >> 32)
+				var_type = VariableSourceType(instr.operands[i] >> 32)
 				index = instr.operands[i] & 0xffffffff
 				identifier = instr.operands[i + 1]
 				i += 1
-				value = function.ILVariable(self.function, var_type, index, identifier)
+				value = function.Variable(self.function, var_type, index, identifier)
 			elif operand_type == "int_list":
 				count = ctypes.c_ulonglong()
 				operand_list = core.BNMediumLevelILGetOperandList(func.handle, self.expr_index, i, count)
@@ -175,10 +175,10 @@ class MediumLevelILInstruction(object):
 				i += 1
 				value = []
 				for j in xrange(count.value / 2):
-					var_type = ILVariableSourceType(operand_list[j * 2] >> 32)
+					var_type = VariableSourceType(operand_list[j * 2] >> 32)
 					index = operand_list[j * 2] & 0xffffffff
 					identifier = operand_list[(j * 2) + 1]
-					value.append(function.ILVariable(self.function, var_type, index, identifier))
+					value.append(function.Variable(self.function, var_type, index, identifier))
 				core.BNMediumLevelILFreeOperandList(operand_list)
 			elif operand_type == "var_ssa_list":
 				count = ctypes.c_ulonglong()
@@ -186,11 +186,11 @@ class MediumLevelILInstruction(object):
 				i += 1
 				value = []
 				for j in xrange(count.value / 3):
-					var_type = ILVariableSourceType(operand_list[j * 3] >> 32)
+					var_type = VariableSourceType(operand_list[j * 3] >> 32)
 					index = operand_list[j * 3] & 0xffffffff
 					identifier = operand_list[(j * 3) + 1]
 					var_index = operand_list[(j * 3) + 2]
-					value.append((function.ILVariable(self.function, var_type, index, identifier), var_index))
+					value.append((function.Variable(self.function, var_type, index, identifier), var_index))
 				core.BNMediumLevelILFreeOperandList(operand_list)
 			elif operand_type == "expr_list":
 				count = ctypes.c_ulonglong()
@@ -295,7 +295,7 @@ class MediumLevelILInstruction(object):
 		return core.BNGetMediumLevelILSSAMemoryIndexAtILInstruction(self.function.handle, self.instr_index)
 
 	def get_ssa_var_possible_values(self, var, index):
-		var_data = core.BNILVariable()
+		var_data = core.BNVariable()
 		var_data.type = var.type
 		var_data.index = var.index
 		var_data.identifier = var.identifier
@@ -304,7 +304,7 @@ class MediumLevelILInstruction(object):
 		return result
 
 	def get_ssa_var_index(self, var):
-		var_data = core.BNILVariable()
+		var_data = core.BNVariable()
 		var_data.type = var.type
 		var_data.index = var.index
 		var_data.identifier = var.identifier
@@ -314,17 +314,17 @@ class MediumLevelILInstruction(object):
 		if isinstance(reg, str):
 			reg = self.function.arch.regs[reg].index
 		result = core.BNGetMediumLevelILVariableForRegisterAtInstruction(self.function.handle, reg, self.instr_index)
-		return function.ILVariable(self.function.source_function, result.type, result.index, result.identifier)
+		return function.Variable(self.function.source_function, result.type, result.index, result.identifier)
 
 	def get_var_for_flag(self, flag):
 		if isinstance(flag, str):
 			flag = self.function.arch.regs[flag].index
 		result = core.BNGetMediumLevelILVariableForFlagAtInstruction(self.function.handle, flag, self.instr_index)
-		return function.ILVariable(self.function.source_function, result.type, result.index, result.identifier)
+		return function.Variable(self.function.source_function, result.type, result.index, result.identifier)
 
 	def get_var_for_stack_location(self, offset):
 		result = core.BNGetMediumLevelILVariableForStackLocationAtInstruction(self.function.handle, offset, self.instr_index)
-		return function.ILVariable(self.function.source_function, result.type, result.index, result.identifier)
+		return function.Variable(self.function.source_function, result.type, result.index, result.identifier)
 
 	def get_reg_value(self, reg):
 		if isinstance(reg, str):
@@ -654,7 +654,7 @@ class MediumLevelILFunction(object):
 		return core.BNGetMediumLevelILNonSSAInstructionIndex(self.handle, instr)
 
 	def get_ssa_var_definition(self, var, index):
-		var_data = core.BNILVariable()
+		var_data = core.BNVariable()
 		var_data.type = var.type
 		var_data.index = var.index
 		var_data.identifier = var.identifier
@@ -671,7 +671,7 @@ class MediumLevelILFunction(object):
 
 	def get_ssa_var_uses(self, var, index):
 		count = ctypes.c_ulonglong()
-		var_data = core.BNILVariable()
+		var_data = core.BNVariable()
 		var_data.type = var.type
 		var_data.index = var.index
 		var_data.identifier = var.identifier
@@ -692,7 +692,7 @@ class MediumLevelILFunction(object):
 		return result
 
 	def get_ssa_var_value(self, var, index):
-		var_data = core.BNILVariable()
+		var_data = core.BNVariable()
 		var_data.type = var.type
 		var_data.index = var.index
 		var_data.identifier = var.identifier