optimizer.py

# -*- coding: utf-8 -*-

from structures import InternalFunction
from lifter import Lifter
from blockstate import ConstantState
from blockstate import CopyState
from blockstate import MemState


from opcodes import *
from baseexecutor import execute_binop, execute_monop
from instructions import MoveInstruction, Instruction

import sys, math
import logging

def apply_peephole_optimizations(func):
	"""
	应用窥视孔优化
	:param func:
	:return:
	"""
	for block in func.graph:
		__order_operands(block)
		__size_one_rewrites(block)
		__size_two_rewrites(block)
		rewrite_free_ptr(block)
		# __sha3_rewrites(block)


def __order_operands(block):
	"""
	排序操作数
	:param block:
	:return:
	"""
	for instruction in block:
		if instruction.opcode not in {"AND", "ADD", "MUL", "OR", "EQ", "XOR"}:
			continue
		reads = instruction.reads
		if isinstance(reads[0], str) != isinstance(reads[1], str):
			if isinstance(reads[0], str):
				instruction.reads.reverse()


def __size_one_rewrites(block):
	instructions = block.get_instructions()
	for i, instruction in enumerate(instructions):
		__remove_self_assign(i, instruction, block)
		__fold_constant(i, instruction, block)
		__rewrite_word_add(instruction)
		__rewrite_shift(i, instruction, block)
		__rewrite_move(i, instruction, block)
		# __rewrite_assert(i, instruction, block)


def __remove_self_assign(i, instruction, block):
	"""
	去掉自我分配
	:param i:
	:param instruction:
	:param block:
	:return:
	"""
	if instruction.opcode == "MOVE" and \
			(instruction.reads[0] == instruction.writes[0]):
		block.set_nop_instruction(i)
	if instruction.opcode == "AND" and instruction.reads[1] == WORD_MASK and \
			(instruction.reads[0] == instruction.writes[1]):
		block.set_nop_instruction(i)


def __fold_constant(i, instruction, block):
	"""
	折叠常数
	:param i:
	:param instruction:
	:param block:
	:return:
	"""
	opcode = instruction.opcode
	constants = instruction.get_constants()
	if not constants:
		return
	inputs = [opcode] + constants
	writes = instruction.writes
	address = instruction.address
	if opcode in bin_ops:
		try:
			value = execute_binop(inputs)
		except ZeroDivisionError:
			return
		new_instruction = MoveInstruction("MOVE", [value], writes, address)
		block.set_instruction(i, new_instruction)
	elif opcode in mono_ops:
		value = execute_monop(inputs)
		new_instruction = MoveInstruction("MOVE", [value], writes, address)
		block.set_instruction(i, new_instruction)


def __rewrite_word_add(instruction):
	"""
	重写ADD
	:param instruction:
	:return:
	"""
	if instruction.opcode != "ADD":
		return
	if instruction.reads[0] != WORD_MASK:
		return
	instruction.opcode = "SUB"
	instruction.reads[0] = instruction.reads[1]
	instruction.reads[1] = 1


def __rewrite_shift(i, instruction, block):
	"""
	重写shift
	:param i:
	:param instruction:
	:param block:
	:return:
	"""
	opcode = instruction.opcode
	address = instruction.address
	if opcode not in {"DIV", "MUL"}:
		return
	if opcode == "DIV":
		num = instruction.reads[1]
		if isinstance(num, str):
			return
		if not (((num & (num - 1)) == 0) and num > 256):
			return
		exp = int(math.log(num, 2))
		new_instruction = \
			Instruction("SR", [instruction.reads[0], exp], instruction.writes, address)
		block.set_instruction(i, new_instruction)
	elif opcode == "MUL":
		num = instruction.reads[0]
		if isinstance(num, str):
			return
		if not (((num & (num - 1)) == 0) and num > 256):
			return
		exp = int(math.log(num, 2))
		new_instruction = \
			Instruction("SL", [instruction.reads[1], exp], instruction.writes, address)
		block.set_instruction(i, new_instruction)


def __rewrite_move(i, instruction, block):
	"""
	重写move
	:param i:
	:param instruction:
	:param block:
	:return:
	"""
	opcode = instruction.opcode
	reads = instruction.reads
	writes = instruction.writes
	address = instruction.address
	if opcode == "DIV" and reads[1] == 1:
		new_instruction = MoveInstruction("MOVE", [reads[0]], writes, address)
		block.set_instruction(i, new_instruction)
	elif opcode == "ADD" and reads[0] == 0:
		new_instruction = MoveInstruction("MOVE", [reads[1]], writes, address)
		block.set_instruction(i, new_instruction)
	elif opcode == "MUL" and reads[0] == 1:
		new_instruction = MoveInstruction("MOVE", [reads[1]], writes, address)
		block.set_instruction(i, new_instruction)
	elif opcode == "AND" and reads[0] == WORD_MASK:
		new_instruction = MoveInstruction("MOVE", [reads[1]], writes, address)
		block.set_instruction(i, new_instruction)
	elif opcode == "SUB" and reads[0] == reads[1]:
		new_instruction = MoveInstruction("MOVE", [0], writes, address)
		block.set_instruction(i, new_instruction)


# def __rewrite_assert(i, instruction, block):
# 	reads = instruction.reads
# 	if instruction.opcode == "JUMPI" and reads[0] == 0:
# 		new_instruction = Instruction("ASSERT", [reads[1]], [], instruction.address)
# 		block.set_instruction(i, new_instruction)


def __size_two_rewrites(block):
	instructions = block.get_instructions()
	for i in range(1, len(instructions)):
		ins_0, ins_1 = instructions[i - 1:i + 1]
		__remove_double_mask(i, ins_0, ins_1, block)
		__remove_address_mask(i, ins_0, ins_1, block)
		__remove_doube_iszero(i, ins_0, ins_1, block)
	# want to remove double iszero first
	for i in range(1, len(instructions)):
		ins_0, ins_1 = instructions[i - 1:i + 1]
		__rewrite_negate_ops(i, ins_0, ins_1, block)


def __remove_double_mask(i, instruction_0, instruction_1, block):
	"""
	去除双重掩码
	:param i:
	:param instruction_0:
	:param instruction_1:
	:param block:
	:return:
	"""
	if instruction_0.opcode != "AND" or \
		instruction_1.opcode != "AND":
		return
	w_0 = instruction_0.writes[0]
	w_1 = instruction_1.writes[0]
	reads_0 = instruction_0.reads
	reads_1 = instruction_1.reads

	if reads_0[0] != reads_1[0] or \
		w_0 != reads_1[1] or \
		w_0 != w_1:
		return
	block.set_nop_instruction(i - 1)
	instruction_1.reads[1] = reads_0[1]


def __remove_address_mask(i, instruction_0, instruction_1, block):
	"""
	去除地址掩码
	:param i:
	:param instruction_0:
	:param instruction_1:
	:param block:
	:return:
	"""
	if instruction_0.opcode not in {"CALLER", "ADDRESS"} or \
		instruction_1.opcode != "AND":
		return
	w_0 = instruction_0.writes[0]
	w_1 = instruction_1.writes[0]
	reads_1 = instruction_1.reads
	if reads_1[0] != ADDRESS_MASK or reads_1[1] != w_0:
		return
	address = instruction_1.address
	instruction_1 = MoveInstruction("MOVE", [w_0], [w_1], address)
	block.set_instruction(i, instruction_1)


def __remove_doube_iszero(i, instruction_0, instruction_1, block):
	if instruction_0.opcode != "ISZERO" or \
		instruction_1.opcode != "ISZERO" or \
		instruction_1.address != instruction_0.address + 1:
		return
	block.set_nop_instruction(i - 1)
	block.set_nop_instruction(i)


def __rewrite_negate_ops(i, instruction_0, instruction_1, block):
	"""
	重写否定操作
	:param i:
	:param instruction_0:
	:param instruction_1:
	:param block:
	:return:
	"""
	opcode = instruction_0.opcode
	if instruction_0.opcode not in negate_ops or \
		instruction_1.opcode != "ISZERO" or \
		instruction_1.address != instruction_0.address + 1:
		return
	block.set_nop_instruction(i - 1)
	opcode = negate_ops[opcode]
	new_instruction = \
		Instruction(opcode, instruction_0.reads, instruction_0.writes, instruction_0.address)
	block.set_instruction(i, new_instruction)


def __sha3_rewrites(block):
	local_memory = MemState()
	instructions = block.get_instructions()
	for index, instruction in enumerate(instructions):
		if instruction.opcode == "SHA3":
			begin, end = instruction.reads
			if begin == 0 and not isinstance(end, str):
				addresses = range(begin, end, 32)
				items = local_memory.lookup_mapping(addresses)
				if len(items) != 0:
					values, indices = zip(*items)
					for i in indices:
						block.set_nop_instruction(i)
					operation = Instruction("SHA3R", list(values), instruction.writes, instruction.address)
					block.set_instruction(index, operation)
		local_memory.add_mapping(index, instruction)


def rewrite_free_ptr(block):
	for i, instruction in enumerate(block.get_instructions()):
		opcode = instruction.opcode
		if opcode == "MLOAD" and \
			instruction.reads[0] == 64:
			new_instruction = MoveInstruction("MOVE", ["$m"], instruction.writes, instruction.address)
			block.set_instruction(i, new_instruction)
		elif opcode == "MSTORE" and \
			instruction.reads[0] == 64:
			new_instruction = MoveInstruction("MOVE", [instruction.reads[1]], ["$m"], instruction.address)
			block.set_instruction(i, new_instruction)


def push_down_definition(block):
	instructions = block.get_instructions()
	for i, instruction in enumerate(instructions):
		writes = instruction.writes
		if len(writes) != 1:
			continue
		ui = get_single_use(i, instruction, instructions)
		if ui == -1:
			continue
		if can_reach(instruction, i, ui, instructions):
			block.set_nop_instruction(i)
			instruction.writes[0] = instructions[ui].writes[0]
			block.set_instruction(ui, instruction)


def get_single_use(begin, d, instructions):
	wr = d.writes[0]
	ui, count = -1, 0
	killed = False
	for i, u in enumerate(instructions):
		if i < begin + 1:
			continue
		if u.reads_from(wr):
			ui = i
			count += 1
		if u.writes_to(wr):
			killed = True
			break
	u = instructions[ui]
	if count == 1 and killed and u.opcode == "MOVE":
		return ui
	return -1


def can_reach(d, begin, end, instructions):
	mload = d.opcode == "MLOAD"
	if d.opcode in mem_read_ops | {"SLOAD"} and not mload:
		return False

	targets = set(d.writes + d.reads)
	# print(instructions[begin])
	for u in instructions[begin + 1:end]:
		if mload and u.opcode in mem_write_ops:
			return False
		for r in u.writes:
			if r in targets:
				return False
	return True


class Optimizer(Lifter):
	"""优化器"""
	def __init__(self, binary):
		logging.info("优化器初始化")
		Lifter.__init__(self, binary)
		# return
		self.__debug = False

		for func in self.get_all_functions():
			self.__optimize_function(func)

	def __optimize_function(self, func):
		"""
		优化函数
		:param func:
		:return:
		"""
		logging.info("优化器：优化函数："+'{:#x}'.format(func.signature))
		self.change = True
		while self.change:
			self.change = False
			self.__propagate_constant_values(func)
			self.__propagate_copy_instructions(func)
			apply_peephole_optimizations(func)
			for block in func.graph:
				push_down_definition(block)
			self.__eliminate_dead_instructions(func)

	def __propagate_constant_values(self, func):
		# logging.info("优化器：传播常数值")
		self.__outs = dict()
		self.__compute_constant_states(func)
		self.__apply_constant_propagation(func)

	def __compute_constant_states(self, func):
		# logging.info("优化器：计算常数状态")
		graph = func.graph
		for block in graph:
			block_id = block.get_id()
			self.__outs[block_id] = ConstantState()

		change = True
		while change:
			change = False
			for block in graph:
				block_id = block.get_id()
				new_out = ConstantState()

				for pre in graph.get_predecessor_ids(block_id):
					new_out.join(self.__outs[pre])

				for inst in block:
					new_out.add_mapping(inst)

				if not (self.__outs[block_id] == new_out):
					self.__outs[block_id] = new_out
					change = True

	def __apply_constant_propagation(self, func):
		# logging.info("优化器：应用常数传播")
		graph = func.graph
		for block in graph:
			block_id = block.get_id()

			before = ConstantState()
			for pre in graph.get_predecessor_ids(block_id):
				before.join(self.__outs[pre])

			for instruction in block:
				before.apply_mapping(instruction)
				before.add_mapping(instruction)

	def get_liveness_states(self, func):
		# logging.info("优化器：获得活跃度状态")
		self.__outs = dict()
		self._uses, self._defs = dict(), dict()
		self.__compute_liveness_states(func)
		if isinstance(func, InternalFunction):  # must keep all return registers alive
			assert func.exit_id != -1
			self.__outs[func.exit_id] |= set(func.writes)
		return self.__outs

	def __eliminate_dead_instructions(self, func):
		# logging.info("优化器：消除死亡指令")
		self.__outs = dict()
		self._uses, self._defs = dict(), dict()
		self.__compute_liveness_states(func)
		if isinstance(func, InternalFunction):  # must keep all return registers alive
			assert func.exit_id != -1
			self.__outs[func.exit_id] |= set(func.writes)
		self.__apply_instruction_elimination(func)

	def __compute_liveness_states(self, func):
		#  logging.info("优化器：计算活跃状态")
		graph = func.graph
		for block in graph:
			block_id = block.get_id()
			self.__compute_use_def(block)
			self.__outs[block_id] = set()

		change = True
		while change:
			change = False
			for block in graph:
				block_id = block.get_id()
				# anything
				new_out = set()
				for suc in graph.get_successor_ids(block_id):
					u, d = self._uses[suc], self._defs[suc]
					# u is used in suc without definition
					# d is defined in suc, kills out[suc]
					new_out |= (u | (self.__outs[suc] - d))
				if self.__outs[block_id] != new_out:
					self.__outs[block_id] = new_out
					change = True
		for out in self.__outs.values():
			out.add("$m")

	def __compute_use_def(self, block):
		# logging.info("优化器：计算使用def")
		block_id = block.get_id()
		d, u = set(), set()
		for instruction in block:
			for read in instruction.get_read_registers():
				if read not in d:
					u.add(read)
			for write in instruction.writes:
				d.add(write)
		self._defs[block_id] = d
		self._uses[block_id] = u

	def __apply_instruction_elimination(self, func):
		# logging.info("优化器：应用指令消除")
		# change = False
		graph = func.graph
		for block in graph:
			block_id = block.get_id()
			out = self.__outs[block_id]

			new_instructions = list()
			for instruction in reversed(block.get_instructions()):
				reads = instruction.get_read_registers()
				writes = instruction.get_write_registers()
				opcode = instruction.opcode

				if opcode in throw_away_ops and len(writes & out) == 0:
					self.change = True
					continue  # do not include
				else:
					out = (out - writes) | reads
					new_instructions.append(instruction)

			new_instructions.reverse()
			block.set_instructions(new_instructions)

	def __propagate_copy_instructions(self, func):
		# logging.info("优化器：传播复制指令")
		self.__outs = dict()
		self.__compute_copy_states(func)
		self.__apply_copy_propagation(func)

	def __compute_copy_states(self, func):
		# logging.info("优化器：计算复制状态")
		graph = func.graph
		for block in graph:
			block_id = block.get_id()
			self.__outs[block_id] = CopyState()

		change = True
		while change:
			change = False
			for block in graph:
				block_id = block.get_id()
				new_out = CopyState()

				for pre in graph.get_predecessor_ids(block_id):
					new_out.join(self.__outs[pre])

				for instruction in block:
					new_out.add_mapping(instruction)

				if not (self.__outs[block_id] == new_out):
					self.__outs[block_id] = new_out
					change = True

	def __apply_copy_propagation(self, func):
		# logging.info("优化器：应用复制传播")
		graph = func.graph
		for block in graph:
			block_id = block.get_id()
			before = CopyState()

			for pre in graph.get_predecessor_ids(block_id):
				before.join(self.__outs[pre])

			# print(before.mapping)
			for instruction in block:
				before.apply_mapping(instruction)
				before.add_mapping(instruction)
		# self.__outs[block_id] = CopyState()

	def visualize_functions(self, step):
		"""
		可视化功能
		:return:
		"""
		for func in self.get_all_functions():
			func.visualize_function(step)


if __name__ == "__main__":
	input_file = open(sys.argv[1])
	line = input_file.readline().strip()
	if " " in line:
		line = line.split(" ")[1]
	input_file.close()
	a = Optimizer(line)
	if "-d" in sys.argv:
		a.debug_functions()
	if "-v" in sys.argv:
		a.visualize_functions("3.Optimizer.")