#!/usr/bin/python

#
#  Copyright (c) 2014, ARM Limited. All rights reserved.
#
#  This program and the accompanying materials
#  are licensed and made available under the terms and conditions of the BSD License
#  which accompanies this distribution.  The full text of the license may be found at
#  http://opensource.org/licenses/bsd-license.php
#
#  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
#  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
#

import getopt
import operator
import os
import pickle
import sys
from sys import argv
from cStringIO import StringIO

modules = {}
functions = {}
functions_addr = {}

def usage():
	print "-t,--trace: Location of the Trace file"
	print "-s,--symbols: Location of the symbols and modules"

def get_address_from_string(address):
	return int(address.strip("S:").strip("N:").strip("EL2:").strip("EL1:"), 16)

def get_module_from_addr(modules, addr):
	for key,value in modules.items():
		if (value['start'] <= addr) and (addr <= value['end']):
			return key
	return None

def add_cycles_to_function(functions, func_name, addr, cycles):
	if func_name != "<Unknown>":
		# Check if we are still in the previous function
		if add_cycles_to_function.prev_func_name == func_name:
			add_cycles_to_function.prev_entry['cycles'] += cycles
			return (add_cycles_to_function.prev_func_name, add_cycles_to_function.prev_module_name)

		if func_name in functions.keys():
			for module_name, module_value in functions[func_name].iteritems():
				if (module_value['start'] <= addr) and (addr < module_value['end']):
					module_value['cycles'] += cycles

					add_cycles_to_function.prev_func_name   = func_name
					add_cycles_to_function.prev_module_name = module_name
					add_cycles_to_function.prev_entry       = module_value
					return (func_name, module_name)
				elif (module_value['end'] == 0):
					module_value['cycles'] += cycles

					add_cycles_to_function.prev_func_name   = func_name
					add_cycles_to_function.prev_module_name = module_name
					add_cycles_to_function.prev_entry       = module_value
					return (func_name, module_name)

		# Workaround to fix the 'info func' limitation that does not expose the 'static' function
		module_name = get_module_from_addr(modules, addr)
		functions[func_name] = {}
		functions[func_name][module_name] = {}
		functions[func_name][module_name]['start']  = 0
		functions[func_name][module_name]['end']    = 0
		functions[func_name][module_name]['cycles'] = cycles
		functions[func_name][module_name]['count']  = 0

		add_cycles_to_function.prev_func_name   = func_name
		add_cycles_to_function.prev_module_name = module_name
		add_cycles_to_function.prev_entry       = functions[func_name][module_name]
		return (func_name, module_name)
	else:
		# Check if we are still in the previous function
		if (add_cycles_to_function.prev_entry is not None) and (add_cycles_to_function.prev_entry['start'] <= addr) and (addr < add_cycles_to_function.prev_entry['end']):
			add_cycles_to_function.prev_entry['cycles'] += cycles
			return (add_cycles_to_function.prev_func_name, add_cycles_to_function.prev_module_name)

		# Generate the key for the given address
		key = addr & ~0x0FFF

		if key not in functions_addr.keys():
			if 'Unknown' not in functions.keys():
				functions['Unknown'] = {}
			if 'Unknown' not in functions['Unknown'].keys():
				functions['Unknown']['Unknown'] = {}
				functions['Unknown']['Unknown']['cycles'] = 0
				functions['Unknown']['Unknown']['count'] = 0
			functions['Unknown']['Unknown']['cycles'] += cycles

			add_cycles_to_function.prev_func_name = None
			return None

		for func_key, module in functions_addr[key].iteritems():
			for module_key, module_value in module.iteritems():
				if (module_value['start'] <= addr) and (addr < module_value['end']):
					module_value['cycles'] += cycles

					# In case o <Unknown> we prefer to fallback on the direct search
					add_cycles_to_function.prev_func_name   = func_key
					add_cycles_to_function.prev_module_name = module_key
					add_cycles_to_function.prev_entry       = module_value
					return (func_key, module_key)

	print "Warning: Function %s @ 0x%x not found" % (func_name, addr)

	add_cycles_to_function.prev_func_name = None
	return None

# Static variables for the previous function
add_cycles_to_function.prev_func_name = None
add_cycles_to_function.prev_entry     = None

def trace_read():
	global trace_process
	line = trace.readline()
	trace_process += len(line)
	return line

#
# Parse arguments
#
trace_name = None
symbols_file = None

opts,args = getopt.getopt(sys.argv[1:], "ht:vs:v", ["help","trace=","symbols="])
if (opts is None) or (not opts):
	usage()
	sys.exit()

for o,a in opts:
    if o in ("-h","--help"):
        usage()
        sys.exit()
    elif o in ("-t","--trace"):
        trace_name = a
    elif o in ("-s","--symbols"):
        symbols_file = a
    else:
        assert False, "Unhandled option (%s)" % o

#
# We try first to see if we run the script from DS-5
#
try:
	from arm_ds.debugger_v1 import Debugger
	from arm_ds.debugger_v1 import DebugException

	# Debugger object for accessing the debugger
	debugger = Debugger()

	# Initialisation commands
	ec = debugger.getExecutionContext(0)
	ec.getExecutionService().stop()
	ec.getExecutionService().waitForStop()
	# in case the execution context reference is out of date
	ec = debugger.getExecutionContext(0)

	#
	# Get the module name and their memory range
	#
	info_file = ec.executeDSCommand("info file")
	info_file_str = StringIO(info_file)

	line = info_file_str.readline().strip('\n')
	while line != '':
		if ("Symbols from" in line):
			# Get the module name from the line 'Symbols from "/home/...."'
			module_name = line.split("\"")[1].split("/")[-1]
			modules[module_name] = {}

			# Look for the text section
			line = info_file_str.readline().strip('\n')
			while (line != '') and ("Symbols from" not in line):
				if ("ER_RO" in line):
					modules[module_name]['start'] = get_address_from_string(line.split()[0])
					modules[module_name]['end']   = get_address_from_string(line.split()[2])
					line = info_file_str.readline().strip('\n')
					break;
				if (".text" in line):
					modules[module_name]['start'] = get_address_from_string(line.split()[0])
					modules[module_name]['end']   = get_address_from_string(line.split()[2])
					line = info_file_str.readline().strip('\n')
					break;
				line = info_file_str.readline().strip('\n')
		line = info_file_str.readline().strip('\n')

	#
	# Get the function name and their memory range
	#
	info_func = ec.executeDSCommand("info func")
	info_func_str = StringIO(info_func)

	# Skip the first line 'Low-level symbols ...'
	line = info_func_str.readline().strip('\n')
	func_prev = None
	while line != '':
		# We ignore all the functions after 'Functions in'
		if ("Functions in " in line):
			line = info_func_str.readline().strip('\n')
			while line != '':
				line = info_func_str.readline().strip('\n')
			line = info_func_str.readline().strip('\n')
			continue

		if ("Low-level symbols" in line):
			# We need to fixup the last function of the module
			if func_prev is not None:
				func_prev['end'] = modules[module_name]['end']
				func_prev = None

			line = info_func_str.readline().strip('\n')
			continue

		func_name = line.split()[1]
		func_start = get_address_from_string(line.split()[0])
		module_name = get_module_from_addr(modules, func_start)

		if func_name not in functions.keys():
			functions[func_name] = {}
		functions[func_name][module_name] = {}
		functions[func_name][module_name]['start'] = func_start
		functions[func_name][module_name]['cycles'] = 0
		functions[func_name][module_name]['count'] = 0

		# Set the end address of the previous function
		if func_prev is not None:
			func_prev['end'] = func_start
		func_prev = functions[func_name][module_name]

		line = info_func_str.readline().strip('\n')

	# Fixup the last function
	func_prev['end'] = modules[module_name]['end']

	if symbols_file is not None:
		pickle.dump((modules, functions), open(symbols_file, "w"))
except:
	if symbols_file is None:
		print "Error: Symbols file is required when run out of ARM DS-5"
		sys.exit()

	(modules, functions) = pickle.load(open(symbols_file, "r"))

#
# Build optimized table for the <Unknown> functions
#
functions_addr = {}
for func_key, module in functions.iteritems():
	for module_key, module_value in module.iteritems():
		key = module_value['start'] & ~0x0FFF
		if key not in functions_addr.keys():
			functions_addr[key] = {}
		if func_key not in functions_addr[key].keys():
			functions_addr[key][func_key] = {}
		functions_addr[key][func_key][module_key] = module_value

#
# Process the trace file
#
if trace_name is None:
	sys.exit()

trace = open(trace_name, "r")
trace_size = os.path.getsize(trace_name)
trace_process = 0

# Get the column names from the first line
columns = trace_read().split()
column_addr     = columns.index('Address')
column_cycles   = columns.index('Cycles')
column_function = columns.index('Function')

line = trace_read()
i = 0
prev_callee = None
while line:
	try:
		func_name = line.split('\t')[column_function].strip()
		address   = get_address_from_string(line.split('\t')[column_addr])
		cycles    = int(line.split('\t')[column_cycles])
		callee = add_cycles_to_function(functions, func_name, address, cycles)
		if (prev_callee != None) and (prev_callee != callee):
			functions[prev_callee[0]][prev_callee[1]]['count'] += 1
		prev_callee = callee
	except ValueError:
		pass
	line = trace_read()
	if ((i % 1000000) == 0) and (i != 0):
		percent = (trace_process * 100.00) / trace_size
		print "Processing file ... (%.2f %%)" % (percent)
	i = i + 1

# Fixup the last callee
functions[prev_callee[0]][prev_callee[1]]['count'] += 1

#
# Process results
#
functions_cycles     = {}
all_functions_cycles = {}
total_cycles         = 0

for func_key, module in functions.iteritems():
	for module_key, module_value in module.iteritems():
		key = "%s/%s" % (module_key, func_key)
		functions_cycles[key] = (module_value['cycles'], module_value['count'])
		total_cycles += module_value['cycles']

		if func_key not in all_functions_cycles.keys():
			all_functions_cycles[func_key] = (module_value['cycles'], module_value['count'])
		else:
			all_functions_cycles[func_key] = tuple(map(sum, zip(all_functions_cycles[func_key], (module_value['cycles'], module_value['count']))))

sorted_functions_cycles     = sorted(functions_cycles.iteritems(), key=operator.itemgetter(1), reverse = True)
sorted_all_functions_cycles = sorted(all_functions_cycles.items(), key=operator.itemgetter(1), reverse = True)

print
print "----"
for (key,value) in sorted_functions_cycles[:20]:
	if value[0] != 0:
		print "%s (cycles: %d - %d%%, count: %d)" % (key, value[0], (value[0] * 100) / total_cycles, value[1])
	else:
		break;
print "----"
for (key,value) in sorted_all_functions_cycles[:20]: 
	if value[0] != 0:
		print "%s (cycles: %d - %d%%, count: %d)" % (key, value[0], (value[0] * 100) / total_cycles, value[1])
	else:
		break;