expressions.py

from opcodes import *

from copy import deepcopy

from opcodes import ADDRESS_MASK

TEMP_REGISTER = "$f"


class Expression(object):
	def __init__(self, opcode, reads, writes, address):
		self.opcode = opcode
		self.reads = deepcopy(reads)
		self.writes = deepcopy(writes)

		self.dependencies = dict()

		# maps from index to dependency
		self.address = address

	def set_dependency(self, tar_register, dependency):
		if not isinstance(tar_register, str):
			return

		modified = False
		for index, register in enumerate(self.reads):
			if index in self.dependencies:
				expression = self.dependencies[index]
				modified |= expression.set_dependency(tar_register, dependency)
			elif register == tar_register:
				if dependency.opcode == "MOVE":
					self.__set_move_expression(index, dependency)
				else:
					self.dependencies[index] = deepcopy(dependency)
				modified = True
		return modified

	def __set_move_expression(self, index, dependency):
		result = dependency.get_dependency(0)
		if result:  # skip the move instruction
			self.dependencies[index] = deepcopy(result)
		else:  # use the register directly
			self.reads[index] = dependency.reads[0]

	def reads_register(self, tar_register):
		depends = False
		for index, register in enumerate(self.reads):
			if index in self.dependencies:
				expression = self.dependencies[index]
				depends |= expression.reads_register(tar_register)
			# print(expression)
			elif register == tar_register:
				depends = True
		return depends

	def contains_operations(self, targets):
		contains = self.opcode in targets
		for index, expression in self.dependencies.items():
			contains |= expression.contains_operations(targets)
		return contains

	def writes_to(self, register):
		return register in self.writes

	def get_dependency(self, index):
		if index in self.dependencies:
			return self.dependencies[index]
		return None

	def invalidates(self, other):
		self_writes = self.get_write_registers()
		if len(self_writes & other.get_write_registers()) != 0 \
				or len(self_writes & other.get_read_registers()) != 0:
			return True
		# if other.opcode in and self.opcode in mem_write_ops:
		# 	return True
		if other.contains_operations({"SLOAD"}) and self.contains_operations({"SSTORE"}):
			return True
		if other.contains_operations(mem_read_ops) and self.contains_operations(mem_write_ops):
			return True

		return False

	def get_write_registers(self):
		return set(self.writes)

	def get_read_registers(self):
		registers = set()
		for index, register in enumerate(self.reads):
			if index in self.dependencies:
				expression = self.dependencies[index]
				registers |= expression.get_read_registers()
			elif isinstance(register, str):
				registers.add(register)
		return registers

	def get_read_count(self, target):
		count = 0
		for index, register in enumerate(self.reads):
			if index in self.dependencies:
				expression = self.dependencies[index]
				count += expression.get_read_count(target)
			elif isinstance(register, str) and register == target:
				count += 1
		return count

	# format all dependencies into a string
	def format_dependencies(self, suppress):
		dependencies = []
		for index in range(len(self.reads)):
			dependencies.append(self.format_dependency(index, suppress))
		return "%s(%s)" % (self.opcode, ", ".join(dependencies))

	# format dependency at index into a string
	def format_dependency(self, index, suppress=True):
		if index in self.dependencies:
			return self.dependencies[index].format_dependencies(suppress)
		read = self.reads[index]
		if isinstance(read, str):
			return read
		if read == ADDRESS_MASK:
			return "AD_MASK"
		if read == WORD_MASK:
			return "WD_MASK"
		return "0x%x" % read

	def __str__(self):
		if self.writes:
			return "%s = %s" % (self.writes[0], self.format_dependencies(True))
		else:
			return self.format_dependencies(True)


class MoveExpression(Expression):
	def format_dependencies(self, suppress):
		return self.format_dependency(0, suppress)


# 	def fold_dependencies(self):
# 		dependency_0 = self.get_dependency(0)
# 		if dependency_0:
# 			return dependency_0.fold_dependencies()
# 		return None
#
# 	def get_inferred_type(self):
# 		dependency_0 = self.get_dependency(0)
# 		if dependency_0:
# 			return dependency_0.get_inferred_type()
# 		return RegisterType.raw


class MonoOpExpression(Expression):
	def format_dependencies(self, suppress=False):
		operator = mono_ops[self.opcode]
		if suppress:
			return "%s %s" % (operator, self.format_dependency(0, False))
		return "(%s %s)" % (operator, self.format_dependency(0, False))


class BinOpExpression(Expression):
	def format_dependencies(self, suppress):
		operator = bin_ops[self.opcode]
		if suppress:
			return "%s %s %s" % (self.format_dependency(0, False), operator, self.format_dependency(1, False))
		return "(%s %s %s)" % (self.format_dependency(0, False), operator, self.format_dependency(1, False))


class JumpExpression(Expression):
	def __str__(self):
		return "goto %s" % self.format_dependency(0)


inverted = {
	"LT": ">=",
	"LEQ": ">",
	"GT": "<=",
	"GEQ": "<",
	"NEQ": "=="
}


class JumpIExpression(Expression):
	def __str__(self):
		return "if (%s) goto %s" % (self.format_dependency(1), self.format_dependency(0))

	def get_inverted_condition(self):
		if 1 in self.dependencies:
			dependency = self.dependencies[1]
			opcode = dependency.opcode
			if opcode in {"NOT", "ISZERO"}:
				return "if (%s)" % dependency.format_dependency(0)
			if opcode in inverted:
				operator = inverted[opcode]
				return "if (%s %s %s)" % (dependency.format_dependency(0), operator,
				                          dependency.format_dependency(1))
			# print(opcode)
		return "if (! %s)" % self.format_dependency(1)

	def get_condition(self):
		return "if (%s)" % self.format_dependency(1)


class MstoreExpression(Expression):
	def __str__(self):
		return "M[%s] = %s" % (self.format_dependency(0), self.format_dependency(1))


class SstoreExpression(Expression):
	def __str__(self):
		return "S[%s] = %s" % (self.format_dependency(0), self.format_dependency(1))


class SloadExpression(Expression):
	def format_dependencies(self, suppress):
		return "S[%s]" % self.format_dependency(0)


class MloadExpression(Expression):
	def format_dependencies(self, suppress):
		return "M[%s]" % self.format_dependency(0)


class CallLoadExpression(Expression):
	def format_dependencies(self, suppress):
		return "C[%s]" % self.format_dependency(0)


class SpecialExpression(Expression):
	def format_dependencies(self, suppress):
		return special_ops[self.opcode]


class IntCallExpression(Expression):
	def __str__(self):
		return ",".join(self.writes) + " = " + self.format_dependencies(True)


class PassExpression(Expression):
	def __init__(self, address):
		Expression.__init__(self, "PASS", [], [], address)


class FakeExpression(Expression):
	def format_dependencies(self, suppress):
		operator = fake_ops[self.opcode]
		if suppress:
			return "%s %s %s" % (self.format_dependency(0, False), operator, self.format_dependency(1, False))
		return "(%s %s %s)" % (self.format_dependency(0, False), operator, self.format_dependency(1, False))