#
# File:: binary_stats.rb
# Description:: 
#
# Author:: Derek Ward <derek.ward@rockstarnorth.com>
# Date:: 14th September 2010
#
# Passed in  :-  see OPTIONS ...
# Passed out :-  stderr contains all errors 
#                stdout for all other output. 
# Returns    :-  returns non zero upon detecting any errors 

#-----------------------------------------------------------------------------
# Uses / Requires
#-----------------------------------------------------------------------------
require 'pipeline/config/projects'
require 'pipeline/os/getopt'
require 'pipeline/os/file'
include Pipeline
require 'systemu'
require 'rexml/document'

#-----------------------------------------------------------------------------
# Constants
#-----------------------------------------------------------------------------
OPTIONS = [
	[ "--help", 		"-h", 	OS::Getopt::BOOLEAN, 	"display usage information." ],
	[ '--1stbinary', 	'-b1', 	OS::Getopt::REQUIRED, 	'the 1st binary file OR a linker map file.' ],
	[ '--2ndbinary', 	'-b2', 	OS::Getopt::REQUIRED, 	'the 2nd binary file OR a linker map file.' ],
	[ '--disassembler',	'-d',   OS::Getopt::REQUIRED, 	'the disassembler to use eg. full path to ps3bin.exe - do not specify for default MS linker map file analysis.' ],
	[ '--errorthreshold',	'-e',   OS::Getopt::REQUIRED, 	'number of bytes that would be classed as an error for all sections.' ],
	[ '--warningthreshold',	'-w',   OS::Getopt::REQUIRED, 	'number of bytes that would be classed as an warning for all sections.' ],
	[ '--state',		'-s',   OS::Getopt::REQUIRED, 	'optional cruise control state file.' ],
	[ '--mods',		'-m',   OS::Getopt::REQUIRED, 	'optional cruise control modifications file.' ],
	[ "--show_symbol_deltas","-sym",OS::Getopt::BOOLEAN, 	"display symbol deltas" ],
	
]

INFO			= "[colourise=black]INFO_MSG: "
INFO_BLUE		= "[colourise=blue]INFO_MSG: "
INFO_BLUE_EMAIL		= "[colourise=blue]INFO_EMA: "
INFO_BLUE_WEB		= "[colourise=blue]INFO_WEB: "
INFO_GREEN		= "[colourise=green]INFO_MSG: "
INFO_ORANGE		= "[colourise=orange]INFO_MSG: "

DEF_WARN 		= (30*1024)	# bytes size that would indicate an error if section > than this
DEF_ERR  		= (70*1024)	# bytes size that would indicate an error if section > than this
JUST			= 26		# justification of output

# keys that appear as output from ps3bin -dsi ( update if ever changes )

KEY_TEXT		= "Text (Code)"
KEY_DATA		= "Data"
KEY_RODATA		= "RO-Data"
KEY_BSS			= "BSS (Uninitialised data)"
KEY_TOTAL		= "Total"

KEYS 			= [ KEY_TEXT, KEY_DATA, KEY_RODATA, KEY_BSS, KEY_TOTAL ]

MAP_SEGMENT_REGEX 	= /([\da-f]*):([\da-f]*)\s([\da-f]*)H\s(\S*)\s*(DATA|CODE)$/ 		# 0001:00000000 00000494H .idata$5                DATA
MAP_SYMBOL_REGEX	= /(\d*):([\da-fA-F]*)\s*(\S*)\s([\da-fA-F]*)\s*(.*)/i 			# 0009:004e8808       ?GADGETTYPE_DIRECTIONAL_MIC@@3VatHashString@rage@@B 8394e408     game4_lib:WeaponTypes.obj
PS3_SYMBOL_REGEX 	= /^0[xX]([A-Fa-f0-9]+)\s(\w{1})\s0[xX]([A-Fa-f0-9]+)\s(.*)$/i 		# 0x0000000000000014 D 0x0014 sys_dbg_initialize_ppu_exception_handler

# constants for symbol delta reporting
MIN_SYMBOL_DELTA		= 1024	# defines the signifcant delta worthy of reporting
MAX_SYMBOL_DELTAS_TO_DISPLAY	= 5	# num to display in report ( email )
JUST_SYMBOLS 			= 12 	# justification of output
ERROR_SYMBOL_PREFIX 		= "Error: Symbol " # error channelling...

#
# Get symbol sizes - could be slow!!!!
def get_symbols( disassembler, bin )

	symbols = {}

	cmd = "#{disassembler} #{bin} -dsy -s -c"
	puts "Executing #{cmd}...\n"
	status, stdout, stderr = systemu(cmd)
	puts "\n...Completed"

	$stderr.puts( cmd ) if ( stderr.length > 0 )
	stderr.each { |err| $stderr.puts "Error: #{err}" }

	stdout.each do |out|	
		matchdata = PS3_SYMBOL_REGEX.match(out)
		if (matchdata and matchdata.length == 5)
			type		= $2
			size 		= $3.to_i(16)
			symbol_name 	= $4
			symbols[symbol_name] = { "size" => size, "type" => type }
		end
	end
	
	symbols
end

#
# Find the section this symbol is in 
def get_section(sym, sections)
	sections.each do |section|
		return section	if ( 	sym["section"] == section["section"] and 
					sym["offset"] >= section["start_addr"] and 
					(sym["offset"] <= section["end_addr"] or section["end_addr"] < 0 ) )

	end
	nil
end

#
# "A list of public symbols, with each address (as section:offset), symbol name, flat address, and .obj file where the symbol is defined"
# http://msdn.microsoft.com/en-gb/windows/hardware/gg463125
def get_symbols_in_ms_map_file(bin)

	syms = []
	symbols = {}
	sections = []

	# read the whole file
	File.open(bin) do |file|								
		while (line = file.gets)		
		
			# Add section
			matchdata = MAP_SEGMENT_REGEX.match(line)
			if (matchdata and matchdata.length == 6)								
				section 		= matchdata[1]
				start_addr 		= matchdata[2].to_i(16)
				size			= matchdata[3].to_i(16)
				name			= matchdata[4]
				classname		= matchdata[5]		
				
				# set the end_addr of the previous section
				sections[-1]["end_addr"] = start_addr - 1 if (sections.length > 0)
					
				# add a new section
				sections << { "section" => section, "start_addr" => start_addr, "size" => size, "name" => name, "classname" => classname, "end_addr" => -1 }
				
				next
			end
		
			# Add symbol
			matchdata = MAP_SYMBOL_REGEX.match(line)
			if (matchdata and matchdata.length == 6)										
				section 		= matchdata[1]
				offset 			= matchdata[2].to_i(16)
				symbol_name		= matchdata[3]
				flat_address		= matchdata[4].to_i(16)
				obj			= matchdata[5]

				syms << { "symbol_name" => symbol_name, "section" => section, "offset" => offset, "flat_address" => flat_address, "obj" => obj }
			end						
		end
	end
	
	# transform symbols into symbols hash with sizes ( calculated from previous flat addresses - I do hope this is a sound way of doing this?! )
	prev = nil
	syms.each do |sym|
		if not prev.nil?
			prev_flat_addr = prev["flat_address"]
			curr_flat_addr = sym["flat_address"]
			
			size = 	curr_flat_addr - prev_flat_addr
			
			size = 0 if (size.abs>1000000000) # hack for now - mmm why are big sizes coming though? I hate the mapfile format not easy to parse.
			
			section = get_section(sym, sections)
			type = section ? "#{section["name"]} (#{section["classname"]})" : ""
						
			symbols[prev["symbol_name"]] = { "size" => size, "type" => type }
		end
		
		prev = sym
	end
	
	puts "#{symbols.length} symbols read"
	
	symbols
end			

#
# pretty print a sigle symbol change
def sym_delta_pretty_print(sym_delta, prefix)
	"#{prefix}#{sym_delta['delta'].to_s.rjust(JUST_SYMBOLS)} [#{sym_delta['old_size'].to_s.rjust(JUST_SYMBOLS)} =>#{sym_delta['new_size'].to_s.rjust(JUST_SYMBOLS)} ] #{sym_delta['type']} #{sym_delta['name']}"
end

#
# Show symbol deltas 
def show_symbol_deltas(symbol_sizes, error_count)
	
	#-------------------------------------------------------------------------
	# Show symbol deltas
	symbol_deltas = []

	symbol_sizes[0].each do |key,old|	
		new_size = (symbol_sizes[1][key]) ? symbol_sizes[1][key]['size'] : 0
		delta =  new_size - old['size']	
		symbol_deltas << { "name" => key, "delta" => delta, "old_size" => old['size'], "new_size" => new_size, "type" => old['type']} # changed or deleted symbols
	end

	symbol_sizes[1].each { |key,new| symbol_deltas << { "name" => key, "delta" => new['size'], "old_size" => 0, "new_size" => new['size'], "type" => new['type']} unless symbol_sizes[0][key] } # added symbols						

	symbol_deltas.delete_if { |s| s['delta'].to_i.abs <= MIN_SYMBOL_DELTA }
	symbol_deltas.sort!{|a,b| b['delta'] <=> a['delta'] }

	# growths
	total_significant_growths = 0
	top_growths,bot_growths = [],[]		
	symbol_deltas.each do |sym_delta|
		if sym_delta['delta'] > 0
			top_growths << sym_delta if 	total_significant_growths < MAX_SYMBOL_DELTAS_TO_DISPLAY
			bot_growths << sym_delta unless total_significant_growths < MAX_SYMBOL_DELTAS_TO_DISPLAY
			total_significant_growths += 1
		end
	end

	# contractions
	total_significant_contractions = 0
	top_contractions,bot_contractions = [],[]
	symbol_deltas.reverse.each do |sym_delta|
		if sym_delta['delta'] < 0
			top_contractions << sym_delta if 	total_significant_contractions < MAX_SYMBOL_DELTAS_TO_DISPLAY
			bot_contractions << sym_delta unless 	total_significant_contractions < MAX_SYMBOL_DELTAS_TO_DISPLAY
			total_significant_contractions += 1
		end
	end	

	if ( total_significant_growths > 0 or total_significant_contractions > 0 )

		prefix1 = INFO_BLUE
		prefix2 = INFO_BLUE_WEB
		prefix3 = INFO_BLUE_EMAIL
		if (error_count>0)
			prefix1 = ERROR_SYMBOL_PREFIX
			prefix3 = ERROR_SYMBOL_PREFIX
		end				

		puts "#{prefix1}Significant Symbol Deltas [ >#{MIN_SYMBOL_DELTA}bytes : Showing #{top_growths.length}/#{total_significant_growths} growths, #{top_contractions.length}/#{total_significant_contractions} contractions. ]"	

		prefix = (error_count>0) ? "Error: Symbol " : ""

		# display growths and contractions of symbols
		if (top_growths.length > 0)			
			top_growths.each { |sym_delta| puts sym_delta_pretty_print(sym_delta,prefix1) }
			puts "#{prefix3}< snipped - see webpage for other #{bot_growths.length} significant growths >" if bot_growths.length > 0
			bot_growths.each { |sym_delta| puts sym_delta_pretty_print(sym_delta,prefix2) }
		end

		if (top_contractions.length > 0)
			bot_growths.reverse.each { |sym_delta| puts sym_delta_pretty_print(sym_delta,prefix2) }				
			puts "#{prefix3}< snipped - see webpage for other #{bot_contractions.length} significant contractions >" if bot_contractions.length > 0			
			top_contractions.reverse.each { |sym_delta| puts sym_delta_pretty_print(sym_delta,prefix1) }
		end			
	end		
end




#-----------------------------------------------------------------------------
# Application entry point
#-----------------------------------------------------------------------------

begin
	g_AppName = File::basename( __FILE__, '.rb' )
	g_ProjectName = ''
	g_BranchName = ''
	g_Project = nil
	g_Config = Pipeline::Config.instance()

	#---------------------------------------------------------------------
	# --- PARSER COMMAND LINE                                          --- 
	#---------------------------------------------------------------------
	opts, trailing = OS::Getopt.getopts( OPTIONS )
	if ( opts['help'] )
		puts OS::Getopt.usage( OPTIONS )
		puts ("Press Enter to continue...")
		$stdin.getc( )				
		Process.exit!( 1 )
	end	

	binaries = []	
	binaries << opts['1stbinary'] if (opts['1stbinary'])
	binaries << opts['2ndbinary'] if (opts['2ndbinary'])

	# eg. 'X:\ps3sdk\dev\usr\local\340_001\cell\host-win32\sn\bin\ps3bin.exe'	
	disassembler = opts['disassembler'] ? opts['disassembler'] : nil

	if (disassembler and not File.exist?(disassembler))
		puts("Warning: Cant find file #{disassembler}")
		Process.exit! -1
	end
	
	show_symbol_deltas = opts['show_symbol_deltas']

	def_warn = opts['warningthreshold'] ? opts['warningthreshold'].to_i : DEF_WARN
	def_err = opts['errorthreshold'] ? opts['errorthreshold'].to_i : DEF_ERR

	if (opts['state'])
		state_file = opts['state']
		if (not File.exist?(state_file))
			puts("Warning: State file specified but cannot be found #{state_file}")
		else			
			doc = REXML::Document.new(File.new(state_file))
			previous_build_state = doc.root.elements["Status"].text
			puts "Previous build state was \'#{previous_build_state}\'"
			if previous_build_state == "Failure"
				puts "Warning: #{g_AppName}: Previous build state was \'#{previous_build_state}\' analysis of executable segments will not cause error for this build."
				def_err = 1000000000
			end
		end
	end

	# DW - NOT USED! - just an idea for now...
	#num_changes = 1
	#
	#if (opts['mods'])
	#	mods_file = opts['mods']
	#	if (not File.exist?(mods_file))
	#		puts("Warning: Modifications file specified but cannot be found #{mods_file}")
	#	else			
	#		doc = REXML::Document.new(File.new(mods_file))
	#		changes = doc.elements.to_a( "//ArrayOfModification/Modification/ChangeNumber" )
	#		num_changes = changes.uniq.length
	#	end
	#end
	# END DW - NOT USED! - just an idea for now...

	# thresholds for sections - configure as required.
	warn_threshold 		= { KEY_TEXT => def_warn, KEY_DATA => def_warn, KEY_RODATA => def_warn, KEY_BSS => def_warn, KEY_TOTAL => def_warn }
	error_threshold 	= { KEY_TEXT => def_err,  KEY_DATA => def_err,  KEY_RODATA => def_err,  KEY_BSS => def_err,  KEY_TOTAL => def_err  }
	fabulous_threshold	= -def_err
	
	error_count, warning_count = 0,0
	
	header_str = ''	
		
	stats  = []
	
	symbol_sizes = []
	
	binaries.each_with_index do |bin, idx|	
		
		puts "Previous\n========" if (idx==0)
		puts "Current\n========"  if (idx==1)

		if (not File.exist?(bin))
			puts("Warning: #{g_AppName}: Cant find file #{bin} no binary/map file analysis cant be done.")
			Process.exit! 0
		end
	
		if (disassembler)
			cmd = "#{disassembler} #{bin} -dsi"
			puts "Executing #{cmd}...\n"
			status, stdout, stderr = systemu(cmd)
			puts "\n...Completed"

			$stderr.puts( cmd ) if ( stderr.length > 0 )
		
			stderr.each do |err|
				error_count += 1
				$stderr.puts "Error: #{err}"
			end
		
			stdout.each_with_index do |out, i|									
				if (i==1)
					vals = out.split
					
					hash = {}
					KEYS.each_with_index do |key,i|
						hash[key] = vals[i]
					end
					stats << hash
				end
			end
						
			symbol_sizes[idx] = get_symbols(disassembler,bin) if (show_symbol_deltas)
		else
			# parse map file			
			started_matching = false

			totals = { KEY_TEXT => 0, KEY_DATA => 0, KEY_RODATA => 0, KEY_BSS => 0, KEY_TOTAL => 0 }

			File.open(bin) do |file|								
				while (line = file.gets)
				
					matchdata = MAP_SEGMENT_REGEX.match(line)
        				if (matchdata and matchdata.length == 6)										
						section 		= matchdata[1]
						start_addr 		= matchdata[2]
						size			= matchdata[3]
						name			= matchdata[4]
						classname		= matchdata[5]
						started_matching 	= true
						
						size			= size.to_i(16) # convert from hex
						
						key = KEY_DATA	

						if 	(classname == "CODE") 
							key = KEY_TEXT
						elsif 	(name.include?(".rdata") )
							key = KEY_RODATA
						elsif 	(name.include?(".bss") )
							key = KEY_BSS
						elsif 	(name.include?(".data") )
							key = KEY_DATA
						end
						
						puts "#{g_AppName} #{line.sub("\n","")} => #{key} #{totals[key]} -> #{totals[key]+size}"
						totals[key] 	  += size												

						totals[KEY_TOTAL] += size

					elsif started_matching
						break
					end	
 				end
			end
									
			stats << totals		
			
			
			# A list of public symbols, with each address (as section:offset), symbol name, flat address, and .obj file where the symbol is defined
			symbol_sizes[idx] = get_symbols_in_ms_map_file(bin) if (show_symbol_deltas)
		end	
	end	
			
	header_str 	= "#{INFO_BLUE}        \t"
	prev_str 	= "#{INFO_BLUE}Previous\t"
	curr_str 	= "#{INFO_BLUE}Current \t"
	delta_str 	= "#{INFO_BLUE}Delta   \t"
	status_str 	= "#{INFO_BLUE}Status  \t"
	
	delta = {}
	
	KEYS.each do |key|
		prev_str += stats[0][key].to_s.ljust(JUST)
		curr_str += stats[1][key].to_s.ljust(JUST)
		header_str += key.ljust(JUST)		

		if (stats[0][key].to_i > 0 and stats[1][key].to_i > 0)		
			delta[key] = stats[1][key].to_i - stats[0][key].to_i			
		else
			delta[key] = 0
		end

		delta_str += delta[key].to_s.ljust(JUST)

		if delta[key] > error_threshold[key]
			status_str += "Err >#{error_threshold[key]}".ljust(JUST) 
			error_count += 1
		elsif delta[key] > warn_threshold[key]
			status_str += "Warn >#{warn_threshold[key]}".ljust(JUST) 
			warning_count += 1
		elsif delta[key] < fabulous_threshold
			status_str += "Champion!".ljust(JUST)		
		else
			status_str += "OK".ljust(JUST)	
		end							
	end
		
	if error_count > 0
		amount = delta["Total"] > 1024 ? "#{delta["Total"].to_f/1024.0}K" : "#{delta["Total"]} bytes"
		$stderr.puts "Error: (#{error_count}) : #{g_AppName}.rb : Large size increase in binary."
		$stderr.puts "Error: #{header_str.sub(INFO_BLUE,"")}"
		$stderr.puts "Error: #{prev_str.sub(INFO_BLUE,"")}"
		$stderr.puts "Error: #{curr_str.sub(INFO_BLUE,"")}"
		$stderr.puts "Error: #{delta_str.sub(INFO_BLUE,"")}"
		$stderr.puts "Error: #{status_str.sub(INFO_BLUE,"")}"
				
	elsif warning_count > 0
		puts "Warning: (#{warning_count}) : #{g_AppName}.rb : Modest size increase in binary."	
	else
		puts "No errors or warnings in #{g_AppName}.rb"
	end

	if error_count == 0
		puts header_str	
		puts prev_str
		puts curr_str
		puts delta_str
		puts status_str
	end
	
	show_symbol_deltas(symbol_sizes, error_count) if show_symbol_deltas and symbol_sizes[0] and symbol_sizes[1] and error_count > 0
	
	Process.exit! -1 if error_count > 0
	
	Process.exit! 0

rescue Exception => ex
	$stderr.puts "Error: Unhandled exception: #{ex.message}"
	$stderr.puts "Backtrace:"
	ex.backtrace.each { |m| $stderr.puts "\t#{m}" }
	Process.exit! -1
end