Embedded memory usage analyser
import argparse
import pathlib
from typing import Optional
import glob


class Args:
    def __init__(self) -> None:
        _parser = argparse.ArgumentParser(
            prog="Nesiac",
        )

        _parser.add_argument("-elf_file", type=pathlib.Path)
        _parser.add_argument("-map_file", type=pathlib.Path)
        _parser.add_argument(
            "directory",
            help="The directory to search for map (.map) and elf (.elf, .out) files. "
            "Not necessary if the map and elf files are provided indivudally.",
            nargs="?",
            type=pathlib.Path,
        )
        _parsed_args = _parser.parse_args()
        efile = _parsed_args.elf_file  # type: Optional[pathlib.Path]
        mfile = _parsed_args.map_file  # type: Optional[pathlib.Path]
        dirpath = _parsed_args.directory  # type: Optional[pathlib.Path]
        if not efile and dirpath:
            globelf = glob.glob(str((dirpath / "*.elf").absolute()))
            globout = glob.glob(str((dirpath / "*.out").absolute()))
            glob_both = globelf + globout
            if len(glob_both) != 1:
                _parser.error(
                    "Elf file was not specified, and there is not "
                    "exactly one elf file in the given directory"
                )
            efile = pathlib.Path(glob_both[0])

        if not mfile and dirpath:
            globd = glob.glob(str((dirpath / "*.map").absolute()))
            if len(globd) != 1:
                _parser.error(
                    "Map file was not specified, and there is not "
                    "exactly one map file in the given directory",
                )
            mfile = pathlib.Path(globd[0])

        if not efile or not mfile:
            _parser.error(
                "If the map and elf files are not both specified, a search directory is required",
            )
        else:
            self._elf_file = efile  # type: pathlib.Path
            self._map_file = mfile  # type: pathlib.Path

    def elf_file(self) -> pathlib.Path:
        return self._elf_file

    def map_file(self) -> pathlib.Path:
        return self._map_file

    _map_text = None  # type: str | None

    def map_text(self) -> str:
        self._map_text = (
            self._map_text or open(self.map_file(), "r", encoding="utf-8").read()
        )
        return self._map_text

nesiac/args.py
Eekle/Nesiac
import time
from rich.layout import Layout
from rich.live import Live
from rich.text import Text
from rich.tree import Tree
from rich.table import Table
from rich.console import Group, Console
from typing import Optional
from . import ingest
from .args import Args
import os

if os.name == "nt":
    import msvcrt

    def get_key() -> Optional[int]:
        if msvcrt.kbhit():
            return msvcrt.getch()[0]
        else:
            return None

else:
    import getchlib

    def get_key() -> Optional[int]:
        stkey = getchlib.getkey(False)  # type: Optional[str]
        if stkey and len(stkey) == 1:
            return ord(stkey)
        return None


class InteractiveRegions:
    def __init__(self, reg: list[ingest.RegionWithSections]) -> None:
        self.regions = reg
        self.r_ix = 0
        self.s_ix = 0
        self.object_page = 0
        self.objects_per_page = 25
        self.sorted_by_size = True
        if len(self.selected_region().children) == 0:
            self._next_region_with_children()
        self.sort_by_size()

    def update_objs_per_page(self, con: Console):
        old = self.objects_per_page
        self.objects_per_page = max(con.height - 10, 1)
        if old != self.objects_per_page:
            self.object_page = 0

    def selected_region(self) -> ingest.RegionWithSections:
        return self.regions[self.r_ix]

    def selected_section(self) -> ingest.SectionWithObjects:
        return self.selected_region().children[self.s_ix]

    def set_selection_to(self, target: ingest.SectionWithObjects):
        for search_r, region in enumerate(self.regions):
            for search_s, section in enumerate(region.children):
                if section == target:
                    self.r_ix = search_r
                    self.s_ix = search_s

    def _next_region_with_children(self) -> bool:
        for ix in range(self.r_ix + 1, len(self.regions)):
            if len(self.regions[ix].children):
                self.r_ix = ix
                return True
        return False

    def _prev_region_with_children(self) -> bool:
        for ix in reversed(range(0, self.r_ix)):
            if len(self.regions[ix].children):
                self.r_ix = ix
                return True
        return False

    def next_section(self):
        # Advance to this region's next child if possible
        if (self.s_ix + 1) < len(self.selected_region().children):
            self.s_ix += 1
            self.object_page = 0
        else:
            # Otherwise move to the next region if possible
            if (self.r_ix + 1) < len(self.regions):
                if self._next_region_with_children():
                    self.s_ix = 0
                    self.object_page = 0

    def objects_in_view(self):
        return self.selected_section().children[
            self.object_page
            * self.objects_per_page : (self.object_page + 1)
            * self.objects_per_page
        ]

    def next_object_page(self):
        if ((self.object_page + 1) * self.objects_per_page) < len(
            self.selected_section().children
        ):
            self.object_page += 1

    def prev_object_page(self):
        if self.object_page > 0:
            self.object_page -= 1

    def prev_section(self):
        if self.s_ix > 0:
            self.s_ix -= 1
            self.object_page = 0
        else:
            if self.r_ix > 0:
                if self._prev_region_with_children():
                    self.s_ix = len(self.selected_region().children) - 1
                    self.object_page = 0

    def reg_section_text(self, reg: ingest.RegionWithSections) -> Text:
        fullness = reg.used_mem() / reg.data.length
        if fullness > 0.9:
            pc_colour = "bright_red"
        elif fullness > 0.75:
            pc_colour = "bright_magenta"
        else:
            pc_colour = ""

        return Text.assemble(
            (reg.data.name, ""),
            " ",
            (f"{fullness:.0%}", pc_colour),
            style="on dark_blue" if reg == self.selected_region() else "",
            no_wrap=True,
        )

    def reg_totals_text(self, reg: ingest.RegionWithSections) -> Text:
        used_size = reg.used_mem()
        used_size_str = size_display(used_size)
        total_size_str = size_display(reg.data.length)
        return Text.assemble(
            (f"{used_size/reg.data.length:.0%}".ljust(4), "bright_green"),
            " | ",
            (f"{used_size_str}".rjust(6), "bright_yellow"),
            " of ",
            (f"{total_size_str}", "bright_yellow"),
        )

    def section_text(self, sec: ingest.SectionWithObjects, total_size: int) -> Text:
        return Text.assemble(
            (f"{sec.data.size/total_size:.0%}".ljust(4), "green"),
            " | ",
            (f"{size_display(sec.data.size)}".rjust(6), "yellow"),
            " | ",
            sec.data.name,
            style="on dark_blue" if sec == self.selected_section() else "",
            no_wrap=True,
        )

    def sort_by_size(self) -> None:
        original_selection = self.selected_section()
        self.sorted_by_size = True
        for region in self.regions:
            for section in region.children:
                section.children.sort(key=lambda x: x.size, reverse=True)
            region.children.sort(key=lambda x: x.data.size, reverse=True)
        self.set_selection_to(original_selection)

    def sort_by_addr(self) -> None:
        original_selection = self.selected_section()
        self.sorted_by_size = False

        for region in self.regions:
            if len(region.children) == 0:
                continue
            for section in region.children:
                section.children.sort(key=lambda x: x.addr)
            region.children.sort(key=lambda s: s.data.addr)
        self.set_selection_to(original_selection)


def size_display(numbytes: int) -> str:
    megabyte = 1024 * 1024
    kilobyte = 1024
    if numbytes > megabyte:
        fval = numbytes / megabyte
        symbol = "M"
    elif numbytes > (kilobyte):
        fval = numbytes / kilobyte
        symbol = "K"
    else:
        return f"{numbytes} B"

    fmt_options = [
        f"{int(fval)}",
        f"{fval:.1f}",
        f"{fval:.2f}",
    ]

    chosen_format = max([f for f in fmt_options if len(f) <= 4], key=len)
    return chosen_format + " " + symbol


def cli() -> None:
    o_data = InteractiveRegions(ingest.ingest())

    def region_view(data: InteractiveRegions, compact: bool) -> Group:
        title = Text.assemble(
            str(Args().elf_file().parts[-1]),
            ", ",
            str(Args().map_file().parts[-1]),
            style="bright_green",
        )
        treeitems = []  # type: list[Tree]
        for reg in data.regions:
            region_title = o_data.reg_section_text(reg)
            is_selected = reg == o_data.selected_region()
            if compact and not is_selected:
                region_title.append(f" [{len(reg.children)}]", style="bright_green")
            reg_tree = Tree(region_title, guide_style="bright_black")
            if (not compact) or reg == o_data.selected_region():
                for sec in reg.children:
                    reg_tree.add(o_data.section_text(sec, reg.data.length))
                reg_tree.add(o_data.reg_totals_text(reg))
            treeitems.append(reg_tree)
        return Group(title, *treeitems)

    def obj_view(data: InteractiveRegions) -> Table:
        numpages = int(
            1 + (len(data.selected_section().children) / data.objects_per_page)
        )
        table = Table(
            title=f"{data.selected_section().data.name} (Page {data.object_page + 1}/{numpages})",
            row_styles=["", "dim"],
            title_justify="left",
        )
        table.add_column("Size", justify="right", style="yellow", width=6)
        table.add_column("Address", justify="right", width=10)
        table.add_column("Symbol Name")
        for obj in data.objects_in_view():
            table.add_row(size_display(obj.size), f"{obj.addr:#010x}", obj.name)

        return table

    def whole_thing(data: InteractiveRegions, viewport_height: int) -> Layout:
        grid = Table.grid()
        grid.add_column()
        grid.add_column()
        num_sections = len([s for reg in data.regions for s in reg.children])
        full_tree_height = 1 + 2 * len(data.regions) + num_sections
        compressed = full_tree_height >= viewport_height
        grid.add_row(
            region_view(data, compressed),
            obj_view(data),
        )
        grid.padding = 2
        layout = Layout()

        layout.split_column(
            Layout(name="u"),
            Layout(
                name="l",
                size=1,
            ),
        )
        sortkind = "Size" if o_data.sorted_by_size else "Addr"
        layout["u"].update(grid)
        layout["l"].update(
            Text(
                "  W/S: Move between sections | "
                "E/D: Scroll symbol table | "
                "Esc: Exit | "
                f"Sorted by (Q): {sortkind} ",
                style="black on white",
            ),
        )
        return layout

    with Live("", auto_refresh=False) as live:
        last_size = live.console.size
        o_data.update_objs_per_page(live.console)
        def render():
            live.update(whole_thing(o_data, live.console.height), refresh=True)
        render()
        while True:
            if key := get_key():
                if key == ord("w"):
                    o_data.prev_section()
                    render()
                elif key == ord("s"):
                    o_data.next_section()
                    render()
                elif key == ord("e"):
                    o_data.prev_object_page()
                    render()
                elif key == ord("d"):
                    o_data.next_object_page()
                    render()
                elif key == ord("q"):
                    if o_data.sorted_by_size:
                        o_data.sort_by_addr()
                    else:
                        o_data.sort_by_size()
                    render()
                elif key == 27:  # Escape
                    live.update("", refresh=True)
                    exit(0)
            else:
                if live.console.size != last_size:
                    last_size = live.console.size
                    o_data.update_objs_per_page(live.console)
                    render()
                time.sleep(5e-3)


if __name__ == "__main__":
    cli()

nesiac/cli.py
Eekle/Nesiac
from dataclasses import replace
from typing import Optional
from elftools.elf.elffile import ELFFile
from elftools.elf.elffile import Section, Segment

from .parsing import (
    program_headers,
    map_regions,
    section_headers,
    objects,
)
from .args import Args


class SectionWithObjects:
    data: section_headers.SectionHeader
    children: list[objects.ElfObject]

    def __init__(self, sec: section_headers.SectionHeader) -> None:
        self.data = sec
        self.children = []

    def add_object(self, obj: objects.ElfObject):
        self.children.append(obj)

    def into_load_section(self, offset: int) -> "SectionWithObjects":
        other = SectionWithObjects(
            sec=replace(
                self.data, name=self.data.name + " (L)", addr=self.data.addr + offset
            )
        )
        other.children = list(self.children)
        return other


class RegionWithSections:
    data: map_regions.MemoryRegion
    children: list[SectionWithObjects]

    def __init__(self, reg: map_regions.MemoryRegion) -> None:
        self.data = reg
        self.children = []

    def add_section(self, sec: SectionWithObjects):
        self.children.append(sec)

    def used_mem(self) -> int:
        return sum(sec.data.size for sec in self.children)


def ingest() -> list[RegionWithSections]:
    args = Args()
    regions = list(map(RegionWithSections, map_regions.from_text(args.map_text())))

    def into_my_section_type(sec: Section) -> Optional[SectionWithObjects]:
        # Check if the section takes up space in the binary
        if (sec.header.sh_flags & 2) == 0:
            return None
        return SectionWithObjects(
            section_headers.SectionHeader(
                sec.name, sec.header.sh_type, sec.header.sh_addr, sec.header.sh_size
            )
        )

    def into_my_header_type(seg: Segment) -> program_headers.ProgramHeader:
        return program_headers.ProgramHeader(
            load_addr=seg.header.p_paddr,
            v_addr=seg.header.p_vaddr,
            size=seg.header.p_memsz,
            typ=seg.header.p_type,
        )

    with open(args.elf_file(), "rb") as elff:
        elf = ELFFile(elff)

        p_hdrs = list(map(into_my_header_type, elf.iter_segments()))
        sections = [
            sec
            for sec in map(into_my_section_type, elf.iter_sections())
            if sec is not None
        ]

        for p_h in p_hdrs:
            if p_h.load_addr != p_h.v_addr:
                new_sections = []  # type: list[SectionWithObjects]
                for section in sections:
                    if section.data.type == "SHT_PROGBITS":
                        if p_h.v_addr <= section.data.addr < (p_h.v_addr + p_h.size):
                            new_sections.append(
                                section.into_load_section(p_h.load_addr - p_h.v_addr)
                            )
                sections.extend(new_sections)

    objs = objects.from_elf(args.elf_file())

    # objs = objects.from_text(external_cmds.elf_objects(args.elf_file()))

    for obj in objs:
        try:
            next(
                s
                for s in sections
                if (s.data.addr <= obj.addr < (s.data.addr + s.data.size))
            ).add_object(obj)
        except StopIteration:
            # Not in a physical section, which is fine
            pass

    for sec in sections:
        for reg in regions:
            if reg.data.contains_address(sec.data.addr):
                reg.add_section(sec)

    for reg in regions:
        reg.children.sort(key=lambda x: x.data.size, reverse=True)
        for sec in reg.children:
            sec.children.sort(key=lambda x: x.size, reverse=True)

    return regions

nesiac/ingest.py
Eekle/Nesiac
from dataclasses import dataclass
from typing import Optional


@dataclass
class MemoryRegion:
    name: str
    origin: int
    length: int

    def contains_address(self, addr: int) -> bool:
        return self.origin <= addr < (self.origin + self.length)


def from_line(line: str) -> Optional[MemoryRegion]:
    pieces = line.split()
    if len(pieces) < 3:
        return None

    try:
        return MemoryRegion(
            name=pieces[0], origin=int(pieces[1], 16), length=int(pieces[2], 16)
        )
    except ValueError:
        return None


def from_text(text: str) -> list[MemoryRegion]:
    iter_lines = iter(text.splitlines())

    for line in iter_lines:
        if "Memory Configuration" in line:
            break

    next(iter_lines)
    next(iter_lines)

    regions = []  # type: list[MemoryRegion]
    for line in iter_lines:
        if reg := from_line(line):
            if reg.name != "*default*":
                regions.append(reg)
        else:
            break
    return regions

nesiac/parsing/map_regions.py
Eekle/Nesiac
from dataclasses import dataclass
import pathlib
from elftools.elf.elffile import ELFFile, SymbolTableSection
import cpp_demangle
import rust_demangler

@dataclass(frozen=True)
class ElfObject:
    addr: int
    size: int
    name: str

def _demangle(input: str) -> str:
    try:
        rust_d = rust_demangler.demangle(input)
    except Exception:
        rust_d = None
    try:
        cpp_d = cpp_demangle.demangle(input)
    except Exception:
        cpp_d = None
    return rust_d or cpp_d or input


def from_elf(file: pathlib.Path) -> list[ElfObject]:
    outlist = []  # type: list[ElfObject]
    with open(file, "rb") as in_f:
        elf = ELFFile(in_f)
        symtab = elf.get_section_by_name(".symtab")
        if not isinstance(symtab, SymbolTableSection):
            return outlist
        for sym in symtab.iter_symbols():
            if sym.entry.st_info.type in ["STT_OBJECT", "STT_FUNC"]:
                outlist.append(
                    ElfObject(sym.entry.st_value, sym.entry.st_size, _demangle(sym.name))
                )
    return outlist

nesiac/parsing/objects.py
Eekle/Nesiac
from dataclasses import dataclass


@dataclass
class ProgramHeader:
    typ: str
    v_addr: int
    load_addr: int
    size: int

nesiac/parsing/program_headers.py
Eekle/Nesiac
from dataclasses import dataclass


@dataclass(frozen=True)
class SectionHeader:
    name: str
    type: str
    addr: int
    size: int

nesiac/parsing/section_headers.py
Eekle/Nesiac