feat(memmap): add topological memory view

Present memory usage in hierarchical view. This view maps modules to
their respective segments and sections.

Change-Id: I5c374b46738edbc83133441ff3f4268f08cb011d
Signed-off-by: Harrison Mutai <harrison.mutai@arm.com>
diff --git a/docs/tools/memory-layout-tool.rst b/docs/tools/memory-layout-tool.rst
index ff5188e..8874bd7 100644
--- a/docs/tools/memory-layout-tool.rst
+++ b/docs/tools/memory-layout-tool.rst
@@ -147,6 +147,88 @@
 The script relies on symbols in the symbol table to determine the start, end,
 and limit addresses of each bootloader stage.
 
+Memory Tree
+~~~~~~~~~~~
+
+A hierarchical view of the memory layout can be produced by passing the option
+``-t`` or ``--tree`` to the tool. This gives the start, end, and size of each
+module, their ELF segments as well as sections.
+
+.. code:: shell
+
+    $ poetry run memory -t
+    build-path: build/fvp/release
+    name                                       start        end       size
+    bl1                                            0    400c000    400c000
+    ├── 00                                         0       5de0       5de0
+    │   ├── .text                                  0       5000       5000
+    │   └── .rodata                             5000       5de0        de0
+    ├── 01                                   4034000    40344c5        4c5
+    │   └── .data                            4034000    40344c5        4c5
+    ├── 02                                   4034500    4034a00        500
+    │   └── .stacks                          4034500    4034a00        500
+    ├── 04                                   4034a00    4035600        c00
+    │   └── .bss                             4034a00    4035600        c00
+    └── 03                                   4036000    403b000       5000
+        └── .xlat_table                      4036000    403b000       5000
+    bl2                                      4021000    4034000      13000
+    ├── 00                                   4021000    4027000       6000
+    │   ├── .text                            4021000    4026000       5000
+    │   └── .rodata                          4026000    4027000       1000
+    └── 01                                   4027000    402e000       7000
+        ├── .data                            4027000    4027809        809
+        ├── .stacks                          4027840    4027e40        600
+        ├── .bss                             4028000    4028800        800
+        └── .xlat_table                      4029000    402e000       5000
+    bl2u                                     4021000    4034000      13000
+    ├── 00                                   4021000    4025000       4000
+    │   ├── .text                            4021000    4024000       3000
+    │   └── .rodata                          4024000    4025000       1000
+    └── 01                                   4025000    402b000       6000
+        ├── .data                            4025000    4025065         65
+        ├── .stacks                          4025080    4025480        400
+        ├── .bss                             4025600    4025c00        600
+        └── .xlat_table                      4026000    402b000       5000
+    bl31                                     4003000    4040000      3d000
+    ├── 02                                  ffe00000   ffe03000       3000
+    │   └── .el3_tzc_dram                   ffe00000   ffe03000       3000
+    ├── 00                                   4003000    4010000       d000
+    │   └── .text                            4003000    4010000       d000
+    └── 01                                   4010000    4021000      11000
+        ├── .rodata                          4010000    4012000       2000
+        ├── .data                            4012000    401219d        19d
+        ├── .stacks                          40121c0    40161c0       4000
+        ├── .bss                             4016200    4018c00       2a00
+        ├── .xlat_table                      4019000    4020000       7000
+        └── .coherent_ram                    4020000    4021000       1000
+
+
+The granularity of this view can be modified with the ``--depth`` option. For
+instance, if you only require the tree up to the level showing segment data,
+you can specify the depth with:
+
+.. code::
+
+    $ poetry run memory -t --depth 2
+    build-path: build/fvp/release
+    name                          start        end       size
+    bl1                               0    400c000    400c000
+    ├── 00                            0       5df0       5df0
+    ├── 01                      4034000    40344c5        4c5
+    ├── 02                      4034500    4034a00        500
+    ├── 04                      4034a00    4035600        c00
+    └── 03                      4036000    403b000       5000
+    bl2                         4021000    4034000      13000
+    ├── 00                      4021000    4027000       6000
+    └── 01                      4027000    402e000       7000
+    bl2u                        4021000    4034000      13000
+    ├── 00                      4021000    4025000       4000
+    └── 01                      4025000    402b000       6000
+    bl31                        4003000    4040000      3d000
+    ├── 02                     ffe00000   ffe03000       3000
+    ├── 00                      4003000    4010000       d000
+    └── 01                      4010000    4021000      11000
+
 --------------
 
 *Copyright (c) 2023, Arm Limited. All rights reserved.*
diff --git a/tools/memory/memory/buildparser.py b/tools/memory/memory/buildparser.py
index 0e3beaa..c128c36 100755
--- a/tools/memory/memory/buildparser.py
+++ b/tools/memory/memory/buildparser.py
@@ -59,6 +59,18 @@
                 mem_map[k] = mod_mem_map
         return mem_map
 
+    def get_mem_tree_as_dict(self) -> dict:
+        """Returns _tree of modules, segments and segments and their total
+        memory usage."""
+        return {
+            k: {
+                "name": k,
+                **v.get_mod_mem_usage_dict(),
+                **{"children": v.get_seg_map_as_dict()},
+            }
+            for k, v in self._modules.items()
+        }
+
     @property
     def module_names(self):
         """Returns sorted list of module names."""
diff --git a/tools/memory/memory/elfparser.py b/tools/memory/memory/elfparser.py
index b61f328..1bd68b1 100644
--- a/tools/memory/memory/elfparser.py
+++ b/tools/memory/memory/elfparser.py
@@ -5,11 +5,21 @@
 #
 
 import re
+from dataclasses import asdict, dataclass
 from typing import BinaryIO
 
 from elftools.elf.elffile import ELFFile
 
 
+@dataclass(frozen=True)
+class TfaMemObject:
+    name: str
+    start: int
+    end: int
+    size: int
+    children: list
+
+
 class TfaElfParser:
     """A class representing an ELF file built for TF-A.
 
@@ -29,12 +39,74 @@
             for sym in elf.get_section_by_name(".symtab").iter_symbols()
         }
 
+        self.set_segment_section_map(elf.iter_segments(), elf.iter_sections())
         self._memory_layout = self.get_memory_layout_from_symbols()
+        self._start = elf["e_entry"]
+        self._size, self._free = self._get_mem_usage()
+        self._end = self._start + self._size
 
     @property
     def symbols(self):
         return self._symbols.items()
 
+    @staticmethod
+    def tfa_mem_obj_factory(elf_obj, name=None, children=None, segment=False):
+        """Converts a pyelfparser Segment or Section to a TfaMemObject."""
+        # Ensure each segment is provided a name since they aren't in the
+        # program header.
+        assert not (
+            segment and name is None
+        ), "Attempting to make segment without a name"
+
+        if children is None:
+            children = list()
+
+        # Segment and sections header keys have different prefixes.
+        vaddr = "p_vaddr" if segment else "sh_addr"
+        size = "p_memsz" if segment else "sh_size"
+
+        # TODO figure out how to handle free space for sections and segments
+        return TfaMemObject(
+            name if segment else elf_obj.name,
+            elf_obj[vaddr],
+            elf_obj[vaddr] + elf_obj[size],
+            elf_obj[size],
+            [] if not children else children,
+        )
+
+    def _get_mem_usage(self) -> (int, int):
+        """Get total size and free space for this component."""
+        size = free = 0
+
+        # Use information encoded in the segment header if we can't get a
+        # memory configuration.
+        if not self._memory_layout:
+            return sum(s.size for s in self._segments.values()), 0
+
+        for v in self._memory_layout.values():
+            size += v["length"]
+            free += v["start"] + v["length"] - v["end"]
+
+        return size, free
+
+    def set_segment_section_map(self, segments, sections):
+        """Set segment to section mappings."""
+        segments = list(
+            filter(lambda seg: seg["p_type"] == "PT_LOAD", segments)
+        )
+
+        for sec in sections:
+            for n, seg in enumerate(segments):
+                if seg.section_in_segment(sec):
+                    if n not in self._segments.keys():
+                        self._segments[n] = self.tfa_mem_obj_factory(
+                            seg, name=f"{n:#02}", segment=True
+                        )
+
+                    self._segments[n].children.append(
+                        self.tfa_mem_obj_factory(sec)
+                    )
+
     def get_memory_layout_from_symbols(self, expr=None) -> dict:
         """Retrieve information about the memory configuration from the symbol
         table.
@@ -55,6 +127,10 @@
 
         return memory_layout
 
+    def get_seg_map_as_dict(self):
+        """Get a dictionary of segments and their section mappings."""
+        return [asdict(v) for k, v in self._segments.items()]
+
     def get_elf_memory_layout(self):
         """Get the total memory consumed by this module from the memory
         configuration.
@@ -72,3 +148,14 @@
                 "total": attrs["length"],
             }
         return mem_dict
+
+    def get_mod_mem_usage_dict(self):
+        """Get the total memory consumed by the module, this combines the
+        information in the memory configuration.
+        """
+        return {
+            "start": self._start,
+            "end": self._end,
+            "size": self._size,
+            "free": self._free,
+        }
diff --git a/tools/memory/memory/memmap.py b/tools/memory/memory/memmap.py
index e6b66f9..6d6f39d 100755
--- a/tools/memory/memory/memmap.py
+++ b/tools/memory/memory/memmap.py
@@ -43,6 +43,17 @@
     help="Generate a high level view of memory usage by memory types.",
 )
 @click.option(
+    "-t",
+    "--tree",
+    is_flag=True,
+    help="Generate a hierarchical view of the modules, segments and sections.",
+)
+@click.option(
+    "--depth",
+    default=3,
+    help="Generate a virtual address map of important TF symbols.",
+)
+@click.option(
     "-s",
     "--symbols",
     is_flag=True,
@@ -59,8 +70,10 @@
     root: Path,
     platform: str,
     build_type: str,
-    footprint: bool,
+    footprint: str,
+    tree: bool,
     symbols: bool,
+    depth: int,
     width: int,
     d: bool,
 ):
@@ -70,9 +83,14 @@
     parser = TfaBuildParser(build_path)
     printer = TfaPrettyPrinter(columns=width, as_decimal=d)
 
-    if footprint or not symbols:
+    if footprint or not (tree or symbols):
         printer.print_footprint(parser.get_mem_usage_dict())
 
+    if tree:
+        printer.print_mem_tree(
+            parser.get_mem_tree_as_dict(), parser.module_names, depth=depth
+        )
+
     if symbols:
         expr = (
             r"(.*)(TEXT|BSS|RODATA|STACKS|_OPS|PMF|XLAT|GOT|FCONF"
diff --git a/tools/memory/memory/printer.py b/tools/memory/memory/printer.py
index c6019c2..6bc6bff 100755
--- a/tools/memory/memory/printer.py
+++ b/tools/memory/memory/printer.py
@@ -4,6 +4,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+from anytree import RenderTree
+from anytree.importer import DictImporter
 from prettytable import PrettyTable
 
 
@@ -17,6 +19,7 @@
 
     def __init__(self, columns: int = None, as_decimal: bool = False):
         self.term_size = columns if columns and columns > 120 else 120
+        self._tree = None
         self._footprint = None
         self._symbol_map = None
         self.as_decimal = as_decimal
@@ -26,6 +29,10 @@
             fmt = f">{width}x" if not self.as_decimal else f">{width}"
         return [f"{arg:{fmt}}" if fmt else arg for arg in args]
 
+    def format_row(self, leading, *args, width=10, fmt=None):
+        formatted_args = self.format_args(*args, width=width, fmt=fmt)
+        return leading + " ".join(formatted_args)
+
     @staticmethod
     def map_elf_symbol(
         leading: str,
@@ -125,3 +132,29 @@
         self._symbol_map = ["Memory Layout:"]
         self._symbol_map += list(reversed(_symbol_map))
         print("\n".join(self._symbol_map))
+
+    def print_mem_tree(
+        self, mem_map_dict, modules, depth=1, min_pad=12, node_right_pad=12
+    ):
+        # Start column should have some padding between itself and its data
+        # values.
+        anchor = min_pad + node_right_pad * (depth - 1)
+        headers = ["start", "end", "size"]
+
+        self._tree = [
+            (f"{'name':<{anchor}}" + " ".join(f"{arg:>10}" for arg in headers))
+        ]
+
+        for mod in sorted(modules):
+            root = DictImporter().import_(mem_map_dict[mod])
+            for pre, fill, node in RenderTree(root, maxlevel=depth):
+                leading = f"{pre}{node.name}".ljust(anchor)
+                self._tree.append(
+                    self.format_row(
+                        leading,
+                        node.start,
+                        node.end,
+                        node.size,
+                    )
+                )
+        print("\n".join(self._tree), "\n")