pw_bloat: Enable running size reports using memory regions

This adds a function to the bloat module which runs Bloaty using a
temporary bloaty config generated from memory region symbols in an
ELF file, avoiding the need for an external config file.

This function is not yet used; a future change will enable it within
bloat tooling.

Change-Id: Id10a32484fedef10fa6da288ca48cbec6bdbe804
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/108912
Reviewed-by: Brandon Vu <brandonvu@google.com>
Reviewed-by: Ewout van Bekkum <ewout@google.com>
Commit-Queue: Alexei Frolov <frolv@google.com>

pw_bloat/py/pw_bloat/bloat.py

@@ -19,12 +19,15 @@
 import json
 import logging
 import os
+from pathlib import Path
 import subprocess
 import sys
+import tempfile
 from typing import Iterable, Optional
 
 import pw_cli.log
 
+from pw_bloat.bloaty_config import generate_bloaty_config
 from pw_bloat.label import from_bloaty_tsv
 from pw_bloat.label_output import (BloatTableOutput, LineCharset, RstOutput,
                                    AsciiCharset)
@@ -95,6 +98,48 @@ def run_bloaty(
     return subprocess.check_output(cmd)
 
 
+class NoMemoryRegions(Exception):
+    """Exception raised if an ELF does not define any memory region symbols."""
+
+
+def memory_regions_size_report(
+        elf: Path,
+        additional_data_sources: Iterable[str] = (),
+        extra_args: Iterable[str] = (),
+) -> str:
+    """Runs a size report on an ELF file using pw_bloat memory region symbols.
+
+    Arguments:
+        elf: The ELF binary on which to run.
+        additional_data_sources: Optional hierarchical data sources to display
+            following the root memory regions.
+        extra_args: Additional command line arguments forwarded to bloaty.
+
+    Returns:
+        The bloaty TSV output detailing the size report.
+
+    Raises:
+        NoMemoryRegions: The ELF does not define memory region symbols.
+    """
+    with tempfile.NamedTemporaryFile() as bloaty_config:
+        with open(elf.resolve(), "rb") as infile, open(bloaty_config.name,
+                                                       "w") as outfile:
+            result = generate_bloaty_config(infile,
+                                            enable_memoryregions=True,
+                                            enable_utilization=False,
+                                            out_file=outfile)
+
+            if not result.has_memoryregions:
+                raise NoMemoryRegions(elf.name)
+
+        return run_bloaty(
+            str(elf.resolve()),
+            bloaty_config.name,
+            data_sources=('memoryregions', *additional_data_sources),
+            extra_args=extra_args,
+        ).decode('utf-8')
+
+
 def write_file(filename: str, contents: str, out_dir_file: str) -> None:
     path = os.path.join(out_dir_file, filename)
     with open(path, 'w') as output_file:
@@ -156,12 +201,11 @@ def main() -> int:
                                     gn_arg_dict['out_dir'], data_sources,
                                     extra_args)
 
-    default_data_sources = ['segment_names', 'symbols']
+    default_data_sources = ['symbols']
 
     diff_report = ''
     rst_diff_report = ''
     for curr_diff_binary in gn_arg_dict['binaries']:
-
         curr_extra_args = extra_args.copy()
         data_sources = default_data_sources
 
@@ -173,7 +217,7 @@ def main() -> int:
             data_sources = curr_diff_binary['data_sources']
 
         try:
-            single_output_base = run_bloaty(curr_diff_binary["base"],
+            single_output_base = run_bloaty(curr_diff_binary['base'],
                                             curr_diff_binary['bloaty_config'],
                                             data_sources=data_sources,
                                             extra_args=curr_extra_args)
@@ -185,7 +229,7 @@ def main() -> int:
 
         try:
             single_output_target = run_bloaty(
-                curr_diff_binary["target"],
+                curr_diff_binary['target'],
                 curr_diff_binary['bloaty_config'],
                 data_sources=data_sources,
                 extra_args=curr_extra_args)

pw_bloat/py/pw_bloat/bloaty_config.py

@@ -17,7 +17,7 @@
 import logging
 import re
 import sys
-from typing import BinaryIO, Dict, List, Optional, TextIO
+from typing import BinaryIO, Dict, List, NamedTuple, Optional, TextIO
 
 import pw_cli.argument_types
 from elftools.elf import elffile  # type: ignore
@@ -318,8 +318,26 @@ def generate_utilization_data_source() -> str:
     return '\n'.join(output) + '\n'
 
 
-def generate_bloaty_config(elf_file: BinaryIO, enable_memoryregions: bool,
-                           enable_utilization: bool, out_file: TextIO) -> None:
+class BloatyConfigResult(NamedTuple):
+    has_memoryregions: bool
+    has_utilization: bool
+
+
+def generate_bloaty_config(
+    elf_file: BinaryIO,
+    enable_memoryregions: bool,
+    enable_utilization: bool,
+    out_file: TextIO,
+) -> BloatyConfigResult:
+    """Generates a Bloaty config file from symbols within an ELF.
+
+    Returns:
+        Tuple indicating whether a memoryregions data source, a utilization data
+        source, or both were written.
+    """
+
+    result = [False, False]
+
     if enable_memoryregions:
         # Enable the "memoryregions" data_source if the user provided the
         # required pw_bloat specific symbols in their linker script.
@@ -330,10 +348,14 @@ def generate_bloaty_config(elf_file: BinaryIO, enable_memoryregions: bool,
             _LOG.info('memoryregions data_source is provided')
             out_file.write(
                 generate_memoryregions_data_source(segment_to_memory_region))
+            result[0] = True
 
     if enable_utilization:
         _LOG.info('utilization data_source is provided')
         out_file.write(generate_utilization_data_source())
+        result[1] = True
+
+    return BloatyConfigResult(*result)
 
 
 def main() -> int: