import sys
import getopt
import os
import os.path
import re
from subprocess import *

usage = """
SYNOPSYS:

compare_builds [options] filename(s)

DESCRIPTION:

Prints reports about symbol sizes in an executable file. Also, compares symbol differences between two
executable files. Currently implemented only on PS3.

OPTIONS:

-r, --report
	Reports the biggest symbols in specified file. Cannot be used with --compare.
	Default behaviour if neither of --report and --compare is specified.

-c, --compare
	Prints symbol differences between two executable files. Cannot be used with --report.
	
-f, --filter
	Excludes symbols starting with "LS." or "..LNst" from reports and compare

-u, --uniquelocals
	Do NOT merge symbols starting with LS.#### or LX.#### to remove unique number

-D, --disassemble
	Disassemble all functions that increased in size by at least 100 bytes and at least 10%
	
-l, --limit <number>
	Sets the number of entries to print with --report option. Default: 10
	
-b, --nobbs
	Excludes BSS symbols from reports and compares.
	
-t, --notext
	Excludes text symbols from reports and compares.
	
-d, --nodata
	Excludes data symbols from reports and compares.
	
-e, --norodata
	Excludes read-only data symbols from reports and compares.
	
-o, --noother
	Excludes other symbols from reports and compares.

-h, --help
	Prints this help.
"""

class Symbol:

	def __init__(self, name, type, size):
		self.name = name
		self.type = type
		self.size = size
		
	def __lt__(self, op):
		return self.size < op.size
		
report_limit = 10
filter_symbols = False
merge_locals = True
print_text = True
print_bss = True
print_data = True
print_rodata = True
print_other = True
disassemble_growth = False

#####################################################
# Symbols
#####################################################
def get_symbols_ps3(filename):
	ps3sdk_root = os.getenv( 'SCE_PS3_ROOT' )
	
	if ps3sdk_root == None:
		print "WARNING: You don't have the SCE_PS3_ROOT env variable set. The script will assume"
		print "a value of X:/ps3sdk/dev/usr/local/340_001/cell for it. Set your env variable if"
		print "this is incorrect."
		ps3sdk_root = 'X:/ps3sdk/dev/usr/local/340_001/cell'
				
	path = ps3sdk_root + '/host-win32/sn/bin/ps3bin.exe'
	
	if os.path.exists( path ) == False:
		print "ERROR: The script has guessed the ps3bin.exe path as " + path
		print "but it doesn't exist. Please set the value of SCE_PS3_ROOT env variable"
		print "to point to the correct root of the PS3 SDK. Example:"
		print "X:/ps3sdk/dev/usr/local/340_001/cell"
		sys.exit(1)
	
	cmdline = path + ' -dsy -c ' + filename
	process = Popen( cmdline, 0, None, None, PIPE, None )
	outdata, errdata = process.communicate( None )
	
	pattern = re.compile(r'(L[SX])\.[0-9]+\.')
	symbols = {}
	for line in outdata.splitlines():
		
		line = line.strip()
		if len( line ) == 0 or line.startswith( 'U' ):
			continue
			
		startaddr, type, size, name = line.split( None, 3 )
		size = int( size, 16 )
		if size == 0:
			continue

		if merge_locals == True:
			name = pattern.sub('L.', name )
		
		if symbols.has_key( name ):
			symbol = symbols[ name ]
			symbol.size += size
		else:
			symbol = Symbol( name, type, size )
			symbols[ name ] = symbol
	
	return symbols

# EJ: Parsing XBox symbols sucks. Using an executable that will help us out.
def get_symbols_xbox(filename):

	if os.path.exists( filename ) == False:
		print "ERROR:", filename, "doesn't exist!"
		sys.exit(1)

	rs_tools_root = os.getenv( 'RS_TOOLSROOT' )
	
	if rs_tools_root == None:
		print "WARNING: You don't have the RS_TOOLSROOT env variable set. The script will assume"
		print "a value of x:/gta5/tools for it. Set your env variable if this is incorrect."
		ps3sdk_root = 'x:/gta5/tools'
				
	path = rs_tools_root + '/bin/symbolextract.exe'
	
	if os.path.exists( path ) == False:
		print "ERROR: The script has guessed the symbolextract.exe path as " + path
		print "but it doesn't exist."
		sys.exit(1)
	
	# symbolextract.exe -in x:\gta5\build\dev\game_xenon_final.pdb
	cmdline = path + ' -in ' + filename
	process = Popen( cmdline, 0, None, None, PIPE, None )
	outdata, errdata = process.communicate( None )
	
	symbols = {}
	for line in outdata.splitlines():
		
		line = line.strip()
		if len( line ) == 0 or line.startswith( 'U' ):
			continue
			
		type, size, name = line.split( None, 2 )
		size = int( size, 16 )
		if size == 0:
			continue
		
		if symbols.has_key( name ):
			symbol = symbols[ name ]
			symbol.size += size
		else:
			symbol = Symbol( name, type, size )
			symbols[ name ] = symbol
			
	return symbols

def get_symbols_total_size(symbols):
	total_size = 0
	for symbol in symbols:
		total_size += symbol.size
	return total_size

def print_symbols(symbols):
	for symbol in symbols:
	
		if filter_symbols and ( symbol.name.startswith('LS.') or symbol.name.startswith('..LNst') ):
			continue
	
		type = symbol.type.upper()
		typeflags = {
			'B' : ( print_bss, 'bss' ),
			'D' : ( print_data, 'data' ),
			'R' : ( print_rodata, 'r-o data' ),
			'T' : ( print_text, 'text' ) }

		if typeflags.has_key( type ):
			if typeflags[ type ][0] == False:
				continue
			typecaption = typeflags[ type ][1]
		else:
			if print_other == False:
				continue
			
			if len(type) == 1:
				typecaption = 'other'
			else:
				typecaption = type.lower()
		
		print '%10d %-10s %s' %( symbol.size, typecaption, symbol.name )

#####################################################
# Report
#####################################################
def is_ps3(filename):
	index = filename.rfind(".")
	if index >= 0:
		ext = filename[index:len(filename)]
		if ext == ".self":
			return True

	return False

def is_xbox(filename):
	index = filename.rfind(".")
	if index >= 0:
		ext = filename[index:len(filename)]
		if ext == ".pdb" or ext == ".exe":
			return True

	return False

def report(filename):
	if is_ps3(filename):
		symbols = get_symbols_ps3( filename ).values()
	elif is_xbox(filename):
		symbols = get_symbols_xbox( filename ).values()
	else:
		print "Invalid file extension:", ext
		sys.exit(1)

	symbols.sort()
	symbols.reverse()
	print_symbols( symbols[:report_limit] )

def trim_disassembly(filename, entry, toCut, dest):
	cmdline = os.getenv( 'SCE_PS3_ROOT' ) + '/host-win32/sn/bin/ps3bin.exe --concise --disassemble-symbol="' + entry + '" ' + filename
	process = Popen( cmdline, 0, None, None, PIPE, None )
	outdata, errdata = process.communicate( None )
	
	pattern = re.compile(r' 0x[0123456789abcdefABCDEF]{8}')

	for line in outdata.splitlines():
		if line.find('Disassembly of') != -1:
			continue
		elif len(line) > toCut and line[0]=='0' and line[1]=='x':
			# trim address and opcode off left
			line = line[toCut:]
			# replace eight digit hex address (such as a branch target) since it will never match
			line = pattern.sub(" address",line)
			dest.write(line + '\n')
		elif len(line):
			dest.write('**** ' + line + '\n')
	# help the diff program avoid crossing section boundaries
	dest.write('====\n====\n====\n')
	
#####################################################
# Compare
#####################################################
def compare(filename1, filename2):
	exclusive1 = []
	exclusive2 = []
	shared = []

	if is_ps3(filename1) and is_ps3(filename2):
		symbols1 = get_symbols_ps3( filename1 )
		symbols2 = get_symbols_ps3( filename2 )
	elif is_xbox(filename1) and is_xbox(filename2):
		symbols1 = get_symbols_xbox( filename1 )
		symbols2 = get_symbols_xbox( filename2 )
	else:
		print "Invalid comparison:", filename1, "and", filename2
		sys.exit(1)
	
	for name, symbol in symbols1.items():
		if symbols2.has_key( name ):
			delta = symbols2[ name ].size - symbol.size
			if delta <> 0:
				diffsymbol = Symbol( name, symbol.type, delta )
				shared.append( diffsymbol )
		else:
			exclusive1.append( symbol )
			
	for name, symbol in symbols2.items():
		if symbols1.has_key( name ) == False:
			exclusive2.append( symbol )
			
	exclusive1.sort()
	exclusive1.reverse()
	print 'Symbols exclusive to %s' % filename1
	print_symbols( exclusive1 )
	print ''
	
	exclusive2.sort()
	exclusive2.reverse()
	print 'Symbols exclusive to %s' % filename2
	print_symbols( exclusive2 )
	print ''
	
	shared.sort()
	print 'Symbols common to both files (with deltas instead of sizes)'
	print_symbols( shared )

	if disassemble_growth:
		log1 = open(filename1 + '.txt', 'w')
		log2 = open(filename2 + '.txt', 'w')
		for s in shared:
			# only process functions, and only ones that grew by a certain minimum absolute size:
			if s.type.upper()=='T' and s.size > 100:
				# only process functions that grew by a certain percentage as well, so we don't
				# disassemble huge functions that grew by a relatively small amount:
				origSize = symbols1[s.name].size
				newSize = origSize + s.size
				pct = (newSize * 100 / origSize) - 100
				if (pct >= 10):
					log1.write("This function increased by " + str(s.size) + " bytes (" + str(pct) + "%)\n");
					trim_disassembly(filename1,s.name,21,log1)
					log2.write("This function increased by " + str(s.size) + " bytes (" + str(pct) + "%)\n");
					trim_disassembly(filename2,s.name,21,log2)
		log1.close()
		log2.close()
		print 'viscmp ' + filename1 + '.txt ' + filename2 + '.txt'
		print 'Search for **** to jump to next function.'
				
	print ''
	ex1_size = get_symbols_total_size( exclusive1 )
	ex2_size = get_symbols_total_size( exclusive2 )
	shared_size = get_symbols_total_size( shared )
	print 'Total size of symbols exclusive to %s: %d' % ( filename1,  ex1_size )
	print 'Total size of symbols exclusive to %s: %d' % ( filename2, ex2_size )
	print 'Total delta of symbols common to both files: %d' % shared_size
	print 'Total delta between these files: %d' % ( shared_size + ex2_size - ex1_size )

#####################################################
# Main
#####################################################
def main(args):
	global report_limit
	global filter_symbols
	global merge_locals
	global print_bss
	global print_data
	global print_rodata
	global print_text
	global print_other
	global disassemble_growth
	
	do_report = False
	do_compare = False
	
	if len( args ) == 0:
		print usage
		sys.exit(0)
	
	opts, nonopts = getopt.getopt( args, 'rcful:btdeohD', [ 'report', 'compare', 'filter', 'uniquelocals', 'limit=', 'help', 'nobss', 'notext', 'nodata', 'norodata', 'noother', 'disassemble' ] )
	for opt, param in opts:
		if opt in ( '-r', '--report' ):
			do_report = True
		elif opt in ( '-c', '--compare' ):
			do_compare = True
		elif opt in ( '-f', '--filter' ):
			filter_symbols = True
		elif opt in ( '-u', '--uniquelocals' ):
			merge_locals = False
		elif opt in ( '-D', '--disassemble' ):
			disassemble_growth = True
		elif opt in ( '-l', '--limit' ):
			report_limit = int( param )
		elif opt in ( '-b', '--nobss' ):
			print_bss = False
		elif opt in ( '-t', '--notext' ):
			print_text = False
		elif opt in ( '-d', '--nodata' ):
			print_data = False
		elif opt in ( '-e', '--norodata' ):
			print_rodata = False
		elif opt in ( '-o', '--noother' ):
			print_other = False
		elif opt in ( '-h', '--help' ):
			print usage
			sys.exit(0)
	
	if do_report and do_compare:
		print 'Error: you must specify exactly one between --report and --compare'
		sys.exit(1)
			
	if not do_report and not do_compare:
		do_report = True
		
	if do_report:
		if len( nonopts ) <> 1:
			print 'Error: with --report you must specify exactly one input file'
			sys.exit(1)
		report( nonopts[0] )

	elif do_compare:
		if len( nonopts ) <> 2:
			print 'Error: with --compare you must specify exactly two input files'
			sys.exit(1)
			
		compare( nonopts[0], nonopts[1] )

if __name__ == '__main__':
	main( sys.argv[1:] )