#!/usr/bin/python3
#coding:utf-8
#
# Copyright (c) 2021-2022 Amlogic, Inc. All rights reserved.
#
# SPDX-License-Identifier: MIT
#


from __future__ import print_function

import sys
import os
import argparse

parser = argparse.ArgumentParser(description='total size of each object file in an ld linker map.')
parser.add_argument('map_file', help="A map file generated by passing -M/--print-map to ld during linking.")
parser.add_argument('--combine', action='store_true',
                    help="All object files in an .a archive or in a directory are combined")
args = parser.parse_args()

class SectionSize:
    text = 0
    rela_text = 0
    data = 0  # Including metadata like import tables
    bss = 0
    customize = 0
    system_heap = 0
    system_stack = 0
    rodata = 0
    rom_usage = 0
    ram_usage = 0

    def total(self):
        return self.text + self.data + self.bss

    def add_text_with_rodata(self, section, size):
        if section.startswith('.text') or section.startswith('.rodata'):
            self.text += size
            # Add the .rodata to data in order to match the result of size command with gcc's elf file
            # The elf file generate by riscv then toolchain add the .rodata to .text
            if section.startswith('.rodata'):
                self.rodata += size
        elif section.startswith('.rela.dyn'):
            self.rela_text += size
        elif section.startswith('.bss') or section.startswith('.common') or section.startswith('.sbss'):
            self.bss += size
        elif section.startswith('.data'):
            self.data += size
        elif section.startswith('.heap'):
            SectionSize.system_heap += size
        elif section.startswith('.stack'):
            SectionSize.system_stack += size
        else:
            if (size > 0):
                print("customer section:%s, size:%d" % (section, size))
                self.customize += size
        self.rom_usage = (self.text + self.data)
        self.ram_usage = (self.data + self.bss)

    def add_data_with_rodata(self, section, size):
        # Arch is arm64 and arm32
        if section.startswith('.text'):
            self.text += size
        elif section.startswith('.rela.dyn'):
            self.rela_text += size
        elif section.startswith('.bss') or section.startswith('.common') or section.startswith('.sbss'):
            self.bss += size
        elif section.startswith('.data') or section.startswith('.rodata'):
            # Add the .rodata to data in order to match the result of size command with gcc's elf file
            # The elf file generate by arm toolchain then it add the .rodata to .data
            self.data += size
            if section.startswith('.rodata'):
                self.rodata += size
        elif section.startswith('.heap'):
            SectionSize.system_heap += size
        elif section.startswith('.stack'):
            SectionSize.system_stack += size
        else:
            if (size > 0):
                print("customer section:%s, size:%d" % (section, size))
                self.customize += size
        self.rom_usage = (self.text + self.data)
        self.ram_usage = (self.data + self.bss - self.rodata)

    def add_gcc_section(self, section, size):
        if current_section is None:
            return;
        if section.startswith('.comment'):
            return
        if section.startswith('.debug'):
            return
        if section.startswith('.ARM.attributes'):
            return

        if os.getenv('ARCH') == "riscv":
            self.add_text_with_rodata(section, size)
        else:
            self.add_data_with_rodata(section, size)


def print_codesize_module_text_with_rodata():
    print('---------------------------------------------------------------------------------------------------')
    col_format = "%-20s\t%-12s\t%-12s\t%-10s%-10s\t%-12s\t%-12s\t%-7s"
    print(col_format % ("module file", "ROM(text+data)", "RAM(data+bss)", ".text", "(.rodata)", ".data", ".bss", "customize"))
    col_format = "%-20s\t%-12s\t%-12s\t%-10s(%-10s)\t%-12s\t%-12s\t%-7s"
    for source in sources:
        size = size_by_source[source]
        print(col_format % (os.path.basename(source)[:20], size.rom_usage, size.ram_usage, size.text, size.rodata, size.data, size.bss, size.customize))
    print('---------------------------------------------------------------------------------------------------')

def print_codesize_module_data_with_rodata():
    print('---------------------------------------------------------------------------------------------------')
    col_format = "%-20s\t%-12s\t%-12s\t%-7s\t%-10s%-10s\t%-12s\t%-7s"
    print(col_format % ("module file", "ROM(text+data)", "RAM(data+bss)", ".text", ".data", "(.rodata)", ".bss", "customize"))
    col_format = "%-20s\t%-12s\t%-12s\t%-7s\t%-10s(%-10s)\t%-12s\t%-7s"
    for source in sources:
        size = size_by_source[source]
        print(col_format % (os.path.basename(source)[:20], size.rom_usage, size.ram_usage, size.text, size.data, size.rodata, size.bss, size.customize))
    print('---------------------------------------------------------------------------------------------------')

def print_codesize_summary_text_with_rodata():
    global sumrom
    global sumram
    global sumcode
    global sumdata
    global sumbss
    global sumcustomize
    global sumrodata
    global sys_mem_usage

    for source in sources:
        size = size_by_source[source]
        sumcode += size.text
        sumdata += size.data
        sumbss  += size.bss
        sumrodata += size.rodata
        sumrom += size.rom_usage
        sumram += size.ram_usage
        sumcustomize += size.customize

    print('---------------------------------------------------------------------------------------------------')
    col_format = "%-5s\t%-12s\t%-12s\t%-12s%-12s\t%-12s\t%-12s\t%-7s\t%-7s\t%-7s"
    print(col_format % ("    ", "ROM(text+data)", "RAM(data+bss)", ".text", "(.rodata)", ".data", ".bss", "cust", "stack", "heap" ))
    col_format = "%-5s\t%-12s\t%-12s\t%-12s(%-12s)\t%-12s\t%-12s\t%-7s\t%-7s\t%-7s"
    print(col_format % ("total", sumrom, sumram + sys_mem_usage, sumcode, sumrodata, sumdata, sumbss + sys_mem_usage, sumcustomize, SectionSize.system_stack, SectionSize.system_heap))
    print('---------------------------------------------------------------------------------------------------')

def print_codesize_summary_data_with_rodata():
    global sumrom
    global sumram
    global sumcode
    global sumdata
    global sumbss
    global sumcustomize
    global sumrodata
    global sys_mem_usage

    for source in sources:
        size = size_by_source[source]
        sumcode += size.text
        sumdata += size.data
        sumbss  += size.bss
        sumrodata += size.rodata
        sumrom += size.rom_usage
        sumram += size.ram_usage
        sumcustomize += size.customize

    print('---------------------------------------------------------------------------------------------------')
    col_format = "%-5s\t%-12s\t%-12s\t%-7s\t%-12s%-12s\t%-12s\t%-7s\t%-7s\t%-7s"
    print(col_format % ("    ", "ROM(text+data)", "RAM(data+bss)", ".text", ".data", "(.rodata)", ".bss", "cust", "stack", "heap" ))
    col_format = "%-5s\t%-12s\t%-12s\t%-7s\t%-12s(%-12s)\t%-12s\t%-7s\t%-7s\t%-7s"
    print(col_format % ("total", sumrom, sumram + sys_mem_usage, sumcode, sumdata, sumrodata, sumbss + sys_mem_usage, sumcustomize, SectionSize.system_stack, SectionSize.system_heap))
    print('---------------------------------------------------------------------------------------------------')

size_by_source = {}
with open(args.map_file) as f:
    arch_toolchain = "GCC"
    toolchain_keyword = "toolchains"
    print("%s toolchain map analyzer" % arch_toolchain)

    lines = iter(f)
    for line in lines:
        # Found the start line
        if line.strip() == "Linker script and memory map":
            break

    current_section = None
    split_line = None
    last_addr = 0
    last_size = 0
    for line in lines:
        line = line.strip('\n')
        # Found the end line and break out
        if "OUTPUT(freertos.elf" in line:
            break

        if split_line:
            # Glue a line that was split in two back together
            if line.startswith(' ' * 16):
                line = split_line + line
            else:  # Shouldn't happen
                print("Warning: discarding line ", split_line)
            split_line = None

        if ('size before relaxing' in line):
            continue
        if line.startswith((".", " .", " *fill*")):
            pieces = line.split(None, 3)  # Don't split paths containing spaces
            pieces_num = len(pieces)

            if line.startswith("."):
                # Note: this line might be wrapped, with the size of the section
                # on the next line, but we ignore the size anyway and will ignore that line
                current_section = pieces[0]
                # XCC text section format
                if (pieces_num == 4):
                    source = pieces[-1]
            elif pieces_num == 1 and len(line) > 14:
                # The ld splits the rest of this line onto the next if the section name is too long
                split_line = line
            elif pieces_num >= 3 and "=" not in pieces and "before" not in pieces:
                if pieces[0] == "*fill*":
                    # Fill use the last source to store the fill align data
                    #source = pieces[0]
                    size = int(pieces[-1], 16)
                    if (pieces[-2] == last_addr):
                        # sub the merged size from the same address for fill data
                        size = size - last_size
                    else:
                        last_size = size
                        last_addr = pieces[-2]
                else:
                    source = pieces[-1]
                    size = int(pieces[-2], 16)
                    if (pieces[-3] == last_addr):
                        # Sub the merged size from the same address for text data
                        size = size - last_size
                    else:
                        last_size = size
                        last_addr = pieces[-3]

                if args.combine:
                    if '.a(' in source:
                        # Path/to/archive.a(object.o)
                        source = source[:source.index('.a(') + 2]
                    elif '.dir' in source:
                        if 'drivers.dir' in source:
                            module = source.split('drivers.dir/')[1].split('/')[0]
                            source = f'driver_{module}'
                        else:
                            source = source[:source.find(".dir")]
                    elif source.endswith('.o'):
                        if toolchain_keyword in source:
                            source = 'toolchain_obj'
                        else:
                            source = os.path.basename(source)

                if source not in size_by_source:
                    size_by_source[source] = SectionSize()
                size_by_source[source].add_gcc_section(current_section, size)

sources = list(size_by_source.keys())
sources.sort(key = lambda x: size_by_source[x].total())
sumrom = sumram = sumcode = sumdata = sumbss = sumcustomize = sumrodata = 0
sys_mem_usage = SectionSize.system_stack  + SectionSize.system_heap
if os.getenv('ARCH') == "riscv":
    print_codesize_module_text_with_rodata()
    print_codesize_summary_text_with_rodata()
else:
    print_codesize_module_data_with_rodata()
    print_codesize_summary_data_with_rodata()