refactor: improve readability of symbol table

Make the symbol table produced by the memory mapping script more
readable. Add a generic interface for interacting with ELF binaries.
This interface enables us to get symbols that provide some insights into
TF-A's memory usage.

Change-Id: I6646f817a1d38d6184b837b78039b7465a533c5c
Signed-off-by: Harrison Mutai <harrison.mutai@arm.com>
diff --git a/tools/memory/__init__.py b/tools/memory/__init__.py
new file mode 100644
index 0000000..0b4c8d3
--- /dev/null
+++ b/tools/memory/__init__.py
@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+#
+# Copyright (c) 2023, Arm Limited. All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
diff --git a/tools/memory/memory/__init__.py b/tools/memory/memory/__init__.py
new file mode 100644
index 0000000..0b4c8d3
--- /dev/null
+++ b/tools/memory/memory/__init__.py
@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+#
+# Copyright (c) 2023, Arm Limited. All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
diff --git a/tools/memory/memory/buildparser.py b/tools/memory/memory/buildparser.py
new file mode 100755
index 0000000..6f467cd
--- /dev/null
+++ b/tools/memory/memory/buildparser.py
@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2023, Arm Limited. All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
+
+import re
+from pathlib import Path
+
+from memory.elfparser import TfaElfParser
+
+
+class TfaBuildParser:
+    """A class for performing analysis on the memory layout of a TF-A build."""
+
+    def __init__(self, path: Path):
+        self._modules = dict()
+        self._path = path
+        self._parse_modules()
+
+    def __getitem__(self, module: str):
+        """Returns an TfaElfParser instance indexed by module."""
+        return self._modules[module]
+
+    def _parse_modules(self):
+        """Parse ELF files in the build path."""
+        for elf_file in self._path.glob("**/*.elf"):
+            module_name = elf_file.name.split("/")[-1].split(".")[0]
+            with open(elf_file, "rb") as file:
+                self._modules[module_name] = TfaElfParser(file)
+
+        if not len(self._modules):
+            raise FileNotFoundError(
+                f"failed to find ELF files in path {self._path}!"
+            )
+
+    @property
+    def symbols(self) -> list:
+        return [
+            (*sym, k) for k, v in self._modules.items() for sym in v.symbols
+        ]
+
+    @staticmethod
+    def filter_symbols(symbols: list, regex: str = None) -> list:
+        """Returns a map of symbols to modules."""
+        regex = r".*" if not regex else regex
+        return sorted(
+            filter(lambda s: re.match(regex, s[0]), symbols),
+            key=lambda s: (-s[1], s[0]),
+            reverse=True,
+        )
+
+    @property
+    def module_names(self):
+        """Returns sorted list of module names."""
+        return sorted(self._modules.keys())
diff --git a/tools/memory/memory/elfparser.py b/tools/memory/memory/elfparser.py
new file mode 100644
index 0000000..3964e6c
--- /dev/null
+++ b/tools/memory/memory/elfparser.py
@@ -0,0 +1,33 @@
+#
+# Copyright (c) 2023, Arm Limited. All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
+
+from typing import BinaryIO
+
+from elftools.elf.elffile import ELFFile
+
+
+class TfaElfParser:
+    """A class representing an ELF file built for TF-A.
+
+    Provides a basic interface for reading the symbol table and other
+    attributes of an ELF file. The constructor accepts a file-like object with
+    the contents an ELF file.
+    """
+
+    def __init__(self, elf_file: BinaryIO):
+        self._segments = {}
+        self._memory_layout = {}
+
+        elf = ELFFile(elf_file)
+
+        self._symbols = {
+            sym.name: sym.entry["st_value"]
+            for sym in elf.get_section_by_name(".symtab").iter_symbols()
+        }
+
+    @property
+    def symbols(self):
+        return self._symbols.items()
diff --git a/tools/memory/memory/memmap.py b/tools/memory/memory/memmap.py
new file mode 100755
index 0000000..7057228
--- /dev/null
+++ b/tools/memory/memory/memmap.py
@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+
+#
+# Copyright (c) 2023, Arm Limited. All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
+
+from pathlib import Path
+
+import click
+from memory.buildparser import TfaBuildParser
+from memory.printer import TfaPrettyPrinter
+
+
+@click.command()
+@click.option(
+    "-r",
+    "--root",
+    type=Path,
+    default=None,
+    help="Root containing build output.",
+)
+@click.option(
+    "-p",
+    "--platform",
+    show_default=True,
+    default="fvp",
+    help="The platform targeted for analysis.",
+)
+@click.option(
+    "-b",
+    "--build-type",
+    default="release",
+    help="The target build type.",
+    type=click.Choice(["debug", "release"], case_sensitive=False),
+)
+@click.option(
+    "-s",
+    "--symbols",
+    is_flag=True,
+    show_default=True,
+    default=True,
+    help="Generate a map of important TF symbols.",
+)
+@click.option("-w", "--width", type=int, envvar="COLUMNS")
+@click.option(
+    "-d",
+    is_flag=True,
+    default=False,
+    help="Display numbers in decimal base.",
+)
+def main(
+    root: Path,
+    platform: str,
+    build_type: str,
+    symbols: bool,
+    width: int,
+    d: bool,
+):
+    build_path = root if root else Path("build/", platform, build_type)
+    click.echo(f"build-path: {build_path.resolve()}")
+
+    parser = TfaBuildParser(build_path)
+    printer = TfaPrettyPrinter(columns=width, as_decimal=d)
+
+    if symbols:
+        expr = (
+            r"(.*)(TEXT|BSS|RODATA|STACKS|_OPS|PMF|XLAT|GOT|FCONF"
+            r"|R.M)(.*)(START|END)__$"
+        )
+        printer.print_symbol_table(
+            parser.filter_symbols(parser.symbols, expr), parser.module_names
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/memory/memory/printer.py b/tools/memory/memory/printer.py
new file mode 100755
index 0000000..11fd7f0
--- /dev/null
+++ b/tools/memory/memory/printer.py
@@ -0,0 +1,93 @@
+#
+# Copyright (c) 2023, Arm Limited. All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
+
+
+class TfaPrettyPrinter:
+    """A class for printing the memory layout of ELF files.
+
+    This class provides interfaces for printing various memory layout views of
+    ELF files in a TF-A build. It can be used to understand how the memory is
+    structured and consumed.
+    """
+
+    def __init__(self, columns: int = None, as_decimal: bool = False):
+        self.term_size = columns if columns and columns > 120 else 120
+        self._symbol_map = None
+        self.as_decimal = as_decimal
+
+    def format_args(self, *args, width=10, fmt=None):
+        if not fmt and type(args[0]) is int:
+            fmt = f">{width}x" if not self.as_decimal else f">{width}"
+        return [f"{arg:{fmt}}" if fmt else arg for arg in args]
+
+    @staticmethod
+    def map_elf_symbol(
+        leading: str,
+        section_name: str,
+        rel_pos: int,
+        columns: int,
+        width: int = None,
+        is_edge: bool = False,
+    ):
+        empty_col = "{:{}{}}"
+
+        # Some symbols are longer than the column width, truncate them until
+        # we find a more elegant way to display them!
+        len_over = len(section_name) - width
+        if len_over > 0:
+            section_name = section_name[len_over:-len_over]
+
+        sec_row = f"+{section_name:-^{width-1}}+"
+        sep, fill = ("+", "-") if is_edge else ("|", "")
+
+        sec_row_l = empty_col.format(sep, fill + "<", width) * rel_pos
+        sec_row_r = empty_col.format(sep, fill + ">", width) * (
+            columns - rel_pos - 1
+        )
+
+        return leading + sec_row_l + sec_row + sec_row_r
+
+    def print_symbol_table(
+        self,
+        symbols: list,
+        modules: list,
+        start: int = 11,
+    ):
+        assert len(symbols), "Empty symbol list!"
+        modules = sorted(modules)
+        col_width = int((self.term_size - start) / len(modules))
+
+        num_fmt = "0=#010x" if not self.as_decimal else ">10"
+
+        _symbol_map = [
+            " " * start
+            + "".join(self.format_args(*modules, fmt=f"^{col_width}"))
+        ]
+        last_addr = None
+
+        for i, (name, addr, mod) in enumerate(symbols):
+            # Do not print out an address twice if two symbols overlap,
+            # for example, at the end of one region and start of another.
+            leading = (
+                f"{addr:{num_fmt}}" + " " if addr != last_addr else " " * start
+            )
+
+            _symbol_map.append(
+                self.map_elf_symbol(
+                    leading,
+                    name,
+                    modules.index(mod),
+                    len(modules),
+                    width=col_width,
+                    is_edge=(not i or i == len(symbols) - 1),
+                )
+            )
+
+            last_addr = addr
+
+        self._symbol_map = ["Memory Layout:"]
+        self._symbol_map += list(reversed(_symbol_map))
+        print("\n".join(self._symbol_map))
diff --git a/tools/memory/print_memory_map.py b/tools/memory/print_memory_map.py
deleted file mode 100755
index ef53f7e..0000000
--- a/tools/memory/print_memory_map.py
+++ /dev/null
@@ -1,102 +0,0 @@
-#!/usr/bin/env python3
-#
-# Copyright (c) 2019-2022, Arm Limited and Contributors. All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-#
-
-import re
-import os
-import sys
-import operator
-
-# List of folder/map to parse
-bl_images = ['bl1', 'bl2', 'bl31']
-
-# List of symbols to search for
-blx_symbols = ['__BL1_RAM_START__', '__BL1_RAM_END__',
-                '__BL2_END__',
-                '__BL31_END__',
-                '__RO_START__', '__RO_END_UNALIGNED__', '__RO_END__',
-                '__TEXT_START__', '__TEXT_END__',
-                '__TEXT_RESIDENT_START__', '__TEXT_RESIDENT_END__',
-                '__RODATA_START__', '__RODATA_END__',
-                '__DATA_START__', '__DATA_END__',
-                '__STACKS_START__', '__STACKS_END__',
-                '__BSS_START__', '__BSS_END__',
-                '__COHERENT_RAM_START__', '__COHERENT_RAM_END__',
-                '__CPU_OPS_START__', '__CPU_OPS_END__',
-                '__FCONF_POPULATOR_START__', '__FCONF_POPULATOR_END__',
-                '__GOT_START__', '__GOT_END__',
-                '__PARSER_LIB_DESCS_START__', '__PARSER_LIB_DESCS_END__',
-                '__PMF_TIMESTAMP_START__', '__PMF_TIMESTAMP_END__',
-                '__PMF_SVC_DESCS_START__', '__PMF_SVC_DESCS_END__',
-                '__RELA_START__', '__RELA_END__',
-                '__RT_SVC_DESCS_START__', '__RT_SVC_DESCS_END__',
-                '__BASE_XLAT_TABLE_START__', '__BASE_XLAT_TABLE_END__',
-                '__XLAT_TABLE_START__', '__XLAT_TABLE_END__',
-               ]
-
-# Regex to extract address from map file
-address_pattern = re.compile(r"\b0x\w*")
-
-# List of found element: [address, symbol, file]
-address_list = []
-
-# Get the directory from command line or use a default one
-inverted_print = True
-if len(sys.argv) >= 2:
-    build_dir = sys.argv[1]
-    if len(sys.argv) >= 3:
-        inverted_print = sys.argv[2] == '0'
-else:
-    build_dir = 'build/fvp/debug'
-
-max_len = max(len(word) for word in blx_symbols) + 2
-if (max_len % 2) != 0:
-    max_len += 1
-
-# Extract all the required symbols from the map files
-for image in bl_images:
-    file_path = os.path.join(build_dir, image, '{}.map'.format(image))
-    if os.path.isfile(file_path):
-        with open (file_path, 'rt') as mapfile:
-            for line in mapfile:
-                for symbol in blx_symbols:
-                    skip_symbol = 0
-                    # Regex to find symbol definition
-                    line_pattern = re.compile(r"\b0x\w*\s*" + symbol + "\s= .")
-                    match = line_pattern.search(line)
-                    if match:
-                        # Extract address from line
-                        match = address_pattern.search(line)
-                        if match:
-                            if '_END__' in symbol:
-                                sym_start = symbol.replace('_END__', '_START__')
-                                if [match.group(0), sym_start, image] in address_list:
-                                    address_list.remove([match.group(0), sym_start, image])
-                                    skip_symbol = 1
-                            if skip_symbol == 0:
-                                address_list.append([match.group(0), symbol, image])
-
-# Sort by address
-address_list.sort(key=operator.itemgetter(0))
-
-# Invert list for lower address at bottom
-if inverted_print:
-    address_list = reversed(address_list)
-
-# Generate memory view
-print(('{:-^%d}' % (max_len * 3 + 20 + 7)).format('Memory Map from: ' + build_dir))
-for address in address_list:
-    if "bl1" in address[2]:
-        print(address[0], ('+{:-^%d}+ |{:^%d}| |{:^%d}|' % (max_len, max_len, max_len)).format(address[1], '', ''))
-    elif "bl2" in address[2]:
-        print(address[0], ('|{:^%d}| +{:-^%d}+ |{:^%d}|' % (max_len, max_len, max_len)).format('', address[1], ''))
-    elif "bl31" in address[2]:
-        print(address[0], ('|{:^%d}| |{:^%d}| +{:-^%d}+' % (max_len, max_len, max_len)).format('', '', address[1]))
-    else:
-        print(address[0], ('|{:^%d}| |{:^%d}| +{:-^%d}+' % (max_len, max_len, max_len)).format('', '', address[1]))
-
-print(('{:^20}{:_^%d}   {:_^%d}   {:_^%d}' % (max_len, max_len, max_len)).format('', '', '', ''))
-print(('{:^20}{:^%d}   {:^%d}   {:^%d}' % (max_len, max_len, max_len)).format('address', 'bl1', 'bl2', 'bl31'))