#!/usr/bin/env python3

from executils import captureStdout

from collections import defaultdict, namedtuple

from os.path import normpath

import re, sys

Symbol = namedtuple('Symbol', ('name', 'typ', 'size', 'source', 'lineNo'))

_reSymbolInfo = re.compile(

	r'^([0-9a-f]+\s)?([0-9a-f]+\s)?\s*([A-Za-z])\s([^\t]+)(\t[^\t]+)?$'

	)

def parseSymbolSize(objectFile):

	text = captureStdout(sys.stderr, 'nm -CSl "%s"' % objectFile)

	if text is not None:

		for line in text.split('\n'):

			if line:

				match = _reSymbolInfo.match(line)

				assert match is not None, line

				addr_, sizeStr, typ, name, originStr = match.groups()

				if sizeStr is None:

					continue

				if typ in 'Bb':

					# Symbols in BSS (uninitialized data section) do not

					# contribute to executable size, so ignore them.

					continue

				if originStr is None:

					source = lineNo = None

				else:

					source, lineNo = originStr.lstrip().rsplit(':', 1)

					source = normpath(source)

					lineNo = int(lineNo)

				yield Symbol(name, typ, int(sizeStr, 16), source, lineNo)

if __name__ == '__main__':

	if len(sys.argv) == 2:

		executable = sys.argv[1]

		# Get symbol information.

		symbolsBySource = defaultdict(list)

		for symbol in parseSymbolSize(executable):

			symbolsBySource[symbol.source].append(symbol)

		# Build directory tree.

		def newDict():

			return defaultdict(newDict)

		dirTree = newDict()

		for source, symbols in symbolsBySource.items():

			if source is None:

				parts = [ '(no source)' ]

			else:

				assert source[0] == '/'

				parts = source[1 : ].split('/')

				parts[0] = '/' + parts[0]

			node = dirTree

			for part in parts[ : -1]:

				node = node[part + '/']

			node[parts[-1]] = symbols

		# Combine branches without forks.

		def compactTree(node):

			names = set(node.keys())

			while names:

				name = names.pop()

				content = node[name]

				if isinstance(content, dict) and len(content) == 1:

					subName, subContent = next(iter(content.items()))

					if isinstance(subContent, dict):

						# A directory containing a single directory.

						del node[name]

						node[name + subName] = subContent

						names.add(name + subName)

			for content in node.values():

				if isinstance(content, dict):

					compactTree(content)

		compactTree(dirTree)

		# Compute size of all nodes in the tree.

		def buildSizeTree(node):

			if isinstance(node, dict):

				newNode = {}

				for name, content in node.items():

					newNode[name] = buildSizeTree(content)

				nodeSize = sum(size for size, subNode in newNode.values())

				return nodeSize, newNode

			else:

				nodeSize = sum(symbol.size for symbol in node)

				return nodeSize, node

		totalSize, sizeTree = buildSizeTree(dirTree)

		# Output.

		def printTree(size, node, indent):

			if isinstance(node, dict):

				for name, (contentSize, content) in sorted(

						node.items(),

						key=lambda item: (-item[1][0], item[0])

						):

					print('%s%8d %s' % (indent, contentSize, name))

					printTree(contentSize, content, indent + '  ')

			else:

				for symbol in sorted(

						node,

						key=lambda symbol: (-symbol.size, symbol.name)

						):

					lineNo = symbol.lineNo

					print('%s%8d %s %s %s' % (

						indent, symbol.size, symbol.typ, symbol.name,

						'' if lineNo is None else '(line %d)' % lineNo

						))

		printTree(totalSize, sizeTree, '')

	else:

		print('Usage: python3 sizestats.py executable', file=sys.stderr)

		sys.exit(2)