Adding new value object to hold disjoint sets

3 files changed, 92 insertions, 68 deletions

diff --git a/python/function.py b/python/function.py
index a507534c..51d61792 100644
--- a/python/function.py
+++ b/python/function.py
@@ -51,34 +51,78 @@ class RegisterValue(object):
 	def __init__(self, arch, value):
 		self.type = RegisterValueType(value.state)
 		if value.state == RegisterValueType.EntryValue:
-			self.reg = arch.get_reg_name(value.reg)
+			self.reg = arch.get_reg_name(value.value)
+		elif value.state == RegisterValueType.ConstantValue:
+			self.value = value.value
+		elif value.state == RegisterValueType.StackFrameOffset:
+			self.offset = value.value
+
+	def __repr__(self):
+		if self.type == RegisterValueType.EntryValue:
+			return "<entry %s>" % self.reg
+		if self.type == RegisterValueType.ConstantValue:
+			return "<const %#x>" % self.value
+		if self.type == RegisterValueType.StackFrameOffset:
+			return "<stack frame offset %#x>" % self.offset
+		if self.type == RegisterValueType.ReturnAddressValue:
+			return "<return address>"
+		return "<undetermined>"
+
+
+class ValueRange(object):
+	def __init__(self, start, end, step):
+		self.start = start
+		self.end = end
+		self.step = step
+
+	def __repr__(self):
+		if self.step == 1:
+			return "<range: %#x to %#x>" % (self.start, self.end)
+		return "<range: %#x to %#x, step %#x>" % (self.start, self.end, self.step)
+
+
+class PossibleValueSet(object):
+	def __init__(self, arch, value):
+		self.type = RegisterValueType(value.state)
+		if value.state == RegisterValueType.EntryValue:
+			self.reg = arch.get_reg_name(value.value)
 		elif value.state == RegisterValueType.ConstantValue:
 			self.value = value.value
 		elif value.state == RegisterValueType.StackFrameOffset:
 			self.offset = value.value
 		elif value.state == RegisterValueType.SignedRangeValue:
 			self.offset = value.value
-			self.start = value.rangeStart
-			self.end = value.rangeEnd
-			self.step = value.rangeStep
-			if self.start & (1 << 63):
-				self.start |= ~((1 << 63) - 1)
-			if self.end & (1 << 63):
-				self.end |= ~((1 << 63) - 1)
+			self.ranges = []
+			for i in xrange(0, value.count):
+				start = value.ranges[i].start
+				end = value.ranges[i].end
+				step = value.ranges[i].step
+				if start & (1 << 63):
+					start |= ~((1 << 63) - 1)
+				if end & (1 << 63):
+					end |= ~((1 << 63) - 1)
+				self.ranges.append(ValueRange(start, end, step))
 		elif value.state == RegisterValueType.UnsignedRangeValue:
 			self.offset = value.value
-			self.start = value.rangeStart
-			self.end = value.rangeEnd
-			self.step = value.rangeStep
+			self.ranges = []
+			for i in xrange(0, value.count):
+				start = value.ranges[i].start
+				end = value.ranges[i].end
+				step = value.ranges[i].step
+				self.ranges.append(ValueRange(start, end, step))
 		elif value.state == RegisterValueType.LookupTableValue:
 			self.table = []
 			self.mapping = {}
-			for i in xrange(0, value.rangeEnd):
+			for i in xrange(0, value.count):
 				from_list = []
 				for j in xrange(0, value.table[i].fromCount):
 					from_list.append(value.table[i].fromValues[j])
 					self.mapping[value.table[i].fromValues[j]] = value.table[i].toValue
 				self.table.append(LookupTableEntry(from_list, value.table[i].toValue))
+		elif (value.state == RegisterValueType.InSetOfValues) or (value.state == RegisterValueType.NotInSetOfValues):
+			self.values = set()
+			for i in xrange(0, value.count):
+				self.values.add(value.valueSet[i])
 
 	def __repr__(self):
 		if self.type == RegisterValueType.EntryValue:
@@ -87,12 +131,16 @@ class RegisterValue(object):
 			return "<const %#x>" % self.value
 		if self.type == RegisterValueType.StackFrameOffset:
 			return "<stack frame offset %#x>" % self.offset
-		if (self.type == RegisterValueType.SignedRangeValue) or (self.type == RegisterValueType.UnsignedRangeValue):
-			if self.step == 1:
-				return "<range: %#x to %#x>" % (self.start, self.end)
-			return "<range: %#x to %#x, step %#x>" % (self.start, self.end, self.step)
+		if self.type == RegisterValueType.SignedRangeValue:
+			return "<signed ranges: %s>" % repr(self.ranges)
+		if self.type == RegisterValueType.UnsignedRangeValue:
+			return "<unsigned ranges: %s>" % repr(self.ranges)
 		if self.type == RegisterValueType.LookupTableValue:
 			return "<table: %s>" % ', '.join([repr(i) for i in self.table])
+		if self.type == RegisterValueType.InSetOfValues:
+			return "<in %s>" % repr(self.values)
+		if self.type == RegisterValueType.NotInSetOfValues:
+			return "<not in %s>" % repr(self.values)
 		if self.type == RegisterValueType.ReturnAddressValue:
 			return "<return address>"
 		return "<undetermined>"
@@ -425,7 +473,6 @@ class Function(object):
 			reg = arch.regs[reg].index
 		value = core.BNGetRegisterValueAtInstruction(self.handle, arch.handle, addr, reg)
 		result = RegisterValue(arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_reg_value_after(self, addr, reg, arch=None):
@@ -447,7 +494,6 @@ class Function(object):
 			reg = arch.regs[reg].index
 		value = core.BNGetRegisterValueAfterInstruction(self.handle, arch.handle, addr, reg)
 		result = RegisterValue(arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_stack_contents_at(self, addr, offset, size, arch=None):
@@ -473,7 +519,6 @@ class Function(object):
 			arch = self.arch
 		value = core.BNGetStackContentsAtInstruction(self.handle, arch.handle, addr, offset, size)
 		result = RegisterValue(arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_stack_contents_after(self, addr, offset, size, arch=None):
@@ -481,7 +526,6 @@ class Function(object):
 			arch = self.arch
 		value = core.BNGetStackContentsAfterInstruction(self.handle, arch.handle, addr, offset, size)
 		result = RegisterValue(arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_parameter_at(self, addr, func_type, i, arch=None):
@@ -491,7 +535,6 @@ class Function(object):
 			func_type = func_type.handle
 		value = core.BNGetParameterValueAtInstruction(self.handle, arch.handle, addr, func_type, i)
 		result = RegisterValue(arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_parameter_at_low_level_il_instruction(self, instr, func_type, i):
@@ -499,7 +542,6 @@ class Function(object):
 			func_type = func_type.handle
 		value = core.BNGetParameterValueAtLowLevelILInstruction(self.handle, instr, func_type, i)
 		result = RegisterValue(self.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_regs_read_by(self, addr, arch=None):
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 74b030d7..a737d95e 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -263,15 +263,14 @@ class LowLevelILInstruction(object):
 		"""Value of expression if constant or a known value (read-only)"""
 		value = core.BNGetLowLevelILExprValue(self.function.handle, self.expr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	@property
 	def possible_values(self):
 		"""Possible values of expression using path-sensitive static data flow analysis (read-only)"""
 		value = core.BNGetLowLevelILPossibleExprValues(self.function.handle, self.expr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_reg_value(self, reg):
@@ -279,7 +278,6 @@ class LowLevelILInstruction(object):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetLowLevelILRegisterValueAtInstruction(self.function.handle, reg, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_reg_value_after(self, reg):
@@ -287,23 +285,22 @@ class LowLevelILInstruction(object):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetLowLevelILRegisterValueAfterInstruction(self.function.handle, reg, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_possible_reg_values(self, reg):
 		if isinstance(reg, str):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetLowLevelILPossibleRegisterValuesAtInstruction(self.function.handle, reg, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_possible_reg_values_after(self, reg):
 		if isinstance(reg, str):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetLowLevelILPossibleRegisterValuesAfterInstruction(self.function.handle, reg, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_flag_value(self, flag):
@@ -311,7 +308,6 @@ class LowLevelILInstruction(object):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetLowLevelILFlagValueAtInstruction(self.function.handle, flag, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_flag_value_after(self, flag):
@@ -319,47 +315,44 @@ class LowLevelILInstruction(object):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetLowLevelILFlagValueAfterInstruction(self.function.handle, flag, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_possible_flag_values(self, flag):
 		if isinstance(flag, str):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetLowLevelILPossibleFlagValuesAtInstruction(self.function.handle, flag, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_possible_flag_values_after(self, flag):
 		if isinstance(flag, str):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetLowLevelILPossibleFlagValuesAfterInstruction(self.function.handle, flag, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_stack_contents(self, offset, size):
 		value = core.BNGetLowLevelILStackContentsAtInstruction(self.function.handle, offset, size, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_stack_contents_after(self, offset, size):
 		value = core.BNGetLowLevelILStackContentsAfterInstruction(self.function.handle, offset, size, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_possible_stack_contents(self, offset, size):
 		value = core.BNGetLowLevelILPossibleStackContentsAtInstruction(self.function.handle, offset, size, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_possible_stack_contents_after(self, offset, size):
 		value = core.BNGetLowLevelILPossibleStackContentsAfterInstruction(self.function.handle, offset, size, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def __setattr__(self, name, value):
@@ -1534,7 +1527,6 @@ class LowLevelILFunction(object):
 			reg = self.arch.regs[reg].index
 		value = core.BNGetLowLevelILSSARegisterValue(self.handle, reg, index)
 		result = function.RegisterValue(self.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_ssa_flag_value(self, flag, index):
@@ -1542,7 +1534,6 @@ class LowLevelILFunction(object):
 			flag = self.arch.get_flag_by_name(flag)
 		value = core.BNGetLowLevelILSSAFlagValue(self.handle, flag, index)
 		result = function.RegisterValue(self.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_mapped_medium_level_il_instruction_index(self, instr):
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 6f482221..951ad591 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -260,15 +260,14 @@ class MediumLevelILInstruction(object):
 		"""Value of expression if constant or a known value (read-only)"""
 		value = core.BNGetMediumLevelILExprValue(self.function.handle, self.expr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	@property
 	def possible_values(self):
 		"""Possible values of expression using path-sensitive static data flow analysis (read-only)"""
 		value = core.BNGetMediumLevelILPossibleExprValues(self.function.handle, self.expr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	@property
@@ -302,7 +301,6 @@ class MediumLevelILInstruction(object):
 		var_data.identifier = var.identifier
 		value = core.BNGetMediumLevelILPossibleSSAVarValues(self.function.handle, var_data, index, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_ssa_var_index(self, var):
@@ -333,7 +331,6 @@ class MediumLevelILInstruction(object):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetMediumLevelILRegisterValueAtInstruction(self.function.handle, reg, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_reg_value_after(self, reg):
@@ -341,23 +338,22 @@ class MediumLevelILInstruction(object):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetMediumLevelILRegisterValueAfterInstruction(self.function.handle, reg, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_possible_reg_values(self, reg):
 		if isinstance(reg, str):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetMediumLevelILPossibleRegisterValuesAtInstruction(self.function.handle, reg, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_possible_reg_values_after(self, reg):
 		if isinstance(reg, str):
 			reg = self.function.arch.regs[reg].index
 		value = core.BNGetMediumLevelILPossibleRegisterValuesAfterInstruction(self.function.handle, reg, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_flag_value(self, flag):
@@ -365,7 +361,6 @@ class MediumLevelILInstruction(object):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetMediumLevelILFlagValueAtInstruction(self.function.handle, flag, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_flag_value_after(self, flag):
@@ -373,47 +368,44 @@ class MediumLevelILInstruction(object):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetMediumLevelILFlagValueAfterInstruction(self.function.handle, flag, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_possible_flag_values(self, flag):
 		if isinstance(flag, str):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetMediumLevelILPossibleFlagValuesAtInstruction(self.function.handle, flag, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_possible_flag_values_after(self, flag):
 		if isinstance(flag, str):
 			flag = self.function.arch.flags[flag].index
 		value = core.BNGetMediumLevelILPossibleFlagValuesAfterInstruction(self.function.handle, flag, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_stack_contents(self, offset, size):
 		value = core.BNGetMediumLevelILStackContentsAtInstruction(self.function.handle, offset, size, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_stack_contents_after(self, offset, size):
 		value = core.BNGetMediumLevelILStackContentsAfterInstruction(self.function.handle, offset, size, self.instr_index)
 		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_possible_stack_contents(self, offset, size):
 		value = core.BNGetMediumLevelILPossibleStackContentsAtInstruction(self.function.handle, offset, size, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_possible_stack_contents_after(self, offset, size):
 		value = core.BNGetMediumLevelILPossibleStackContentsAfterInstruction(self.function.handle, offset, size, self.instr_index)
-		result = function.RegisterValue(self.function.arch, value)
-		core.BNFreeRegisterValue(value)
+		result = function.PossibleValueSet(self.function.arch, value)
+		core.BNFreePossibleValueSet(value)
 		return result
 
 	def get_branch_dependence(self, branch_instr):
@@ -706,7 +698,6 @@ class MediumLevelILFunction(object):
 		var_data.identifier = var.identifier
 		value = core.BNGetMediumLevelILSSAVarValue(self.handle, var_data, index)
 		result = function.RegisterValue(self.arch, value)
-		core.BNFreeRegisterValue(value)
 		return result
 
 	def get_low_level_il_instruction_index(self, instr):