xref: /aosp_15_r20/external/pigweed/pw_bloat/py/pw_bloat/label_output.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2022 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Module for size report ASCII tables from DataSourceMaps."""
15
16import enum
17from typing import (
18    Iterable,
19    Type,
20    NamedTuple,
21    cast,
22)
23
24from pw_bloat.label import DataSourceMap, DiffDataSourceMap, Label
25
26
27class AsciiCharset(enum.Enum):
28    """Set of ASCII characters for drawing tables."""
29
30    TL = '+'
31    TM = '+'
32    TR = '+'
33    ML = '+'
34    MM = '+'
35    MR = '+'
36    BL = '+'
37    BM = '+'
38    BR = '+'
39    V = '|'
40    H = '-'
41    HH = '='
42
43
44class LineCharset(enum.Enum):
45    """Set of line-drawing characters for tables."""
46
47    TL = '┌'
48    TM = '┬'
49    TR = '┐'
50    ML = '├'
51    MM = '┼'
52    MR = '┤'
53    BL = '└'
54    BM = '┴'
55    BR = '┘'
56    V = '│'
57    H = '─'
58    HH = '═'
59
60
61class _Align(enum.Enum):
62    CENTER = 0
63    LEFT = 1
64    RIGHT = 2
65
66
67def get_label_status(curr_label: Label) -> str:
68    if curr_label.is_new():
69        return 'NEW'
70    if curr_label.is_del():
71        return 'DEL'
72    return ''
73
74
75def diff_sign_sizes(size: int, diff_mode: bool) -> str:
76    if diff_mode:
77        size_sign = '+' if size > 0 else ''
78        return f"{size_sign}{size:,}"
79    return f"{size:,}"
80
81
82class BloatTableOutput:
83    """ASCII Table generator from DataSourceMap."""
84
85    _RST_PADDING_WIDTH = 6
86    _DEFAULT_MAX_WIDTH = 80
87
88    class _LabelContent(NamedTuple):
89        name: str
90        size: int
91        label_status: str
92
93    def __init__(
94        self,
95        ds_map: DiffDataSourceMap | DataSourceMap,
96        col_max_width: int = _DEFAULT_MAX_WIDTH,
97        charset: Type[AsciiCharset] | Type[LineCharset] = AsciiCharset,
98        rst_output: bool = False,
99        diff_label: str | None = None,
100    ):
101        self._data_source_map = ds_map
102        self._cs = charset
103        self._total_size = 0
104        col_names = [*self._data_source_map.get_ds_names(), 'sizes']
105        self._diff_mode = False
106        self._diff_label = diff_label
107        if isinstance(self._data_source_map, DiffDataSourceMap):
108            col_names = ['diff', *col_names]
109            self._diff_mode = True
110        self._col_names = col_names
111        self._additional_padding = 0
112        self._ascii_table_rows: list[str] = []
113        self._rst_output = rst_output
114        self._total_divider = self._cs.HH.value
115        if self._rst_output:
116            self._total_divider = self._cs.H.value
117            self._additional_padding = self._RST_PADDING_WIDTH
118
119        self._col_widths = self._generate_col_width(col_max_width)
120
121    def _generate_col_width(self, col_max_width: int) -> list[int]:
122        """Find column width for all data sources and sizes."""
123        max_len_size = 0
124        diff_len_col_width = 0
125
126        col_list = [
127            len(ds_name) for ds_name in self._data_source_map.get_ds_names()
128        ]
129        for curr_label in self._data_source_map.labels():
130            self._total_size += curr_label.size
131            max_len_size = max(
132                len(diff_sign_sizes(self._total_size, self._diff_mode)),
133                len(diff_sign_sizes(curr_label.size, self._diff_mode)),
134                max_len_size,
135            )
136            for index, parent_label in enumerate(
137                [*curr_label.parents, curr_label.name]
138            ):
139                if len(parent_label) > col_max_width:
140                    col_list[index] = col_max_width
141                elif len(parent_label) > col_list[index]:
142                    col_list[index] = len(parent_label)
143
144        diff_same = 0
145        if self._diff_mode:
146            col_list = [len('Total'), *col_list]
147            diff_same = len('(SAME)')
148        col_list.append(max(max_len_size, len('sizes'), diff_same))
149
150        if self._diff_label is not None:
151            sum_all_col_names = sum(col_list)
152            if sum_all_col_names < len(self._diff_label):
153                diff_len_col_width = (
154                    len(self._diff_label) - sum_all_col_names
155                ) // len(self._col_names)
156
157        return [
158            (x + self._additional_padding + diff_len_col_width)
159            for x in col_list
160        ]
161
162    def _diff_label_names(
163        self,
164        old_labels: tuple[_LabelContent, ...] | None,
165        new_labels: tuple[_LabelContent, ...],
166    ) -> tuple[_LabelContent, ...]:
167        """Return difference between arrays of labels."""
168
169        if old_labels is None:
170            return new_labels
171        diff_list = []
172        for new_lb, old_lb in zip(new_labels, old_labels):
173            if (new_lb.name == old_lb.name) and (new_lb.size == old_lb.size):
174                diff_list.append(self._LabelContent('', 0, ''))
175            else:
176                diff_list.append(new_lb)
177
178        return tuple(diff_list)
179
180    def _label_title_row(self) -> list[str]:
181        label_rows = []
182        label_cells = ''
183        divider_cells = ''
184        for width in self._col_widths:
185            label_cells += ' ' * width + ' '
186            divider_cells += (self._cs.H.value * width) + self._cs.H.value
187        if self._diff_label is not None:
188            label_cells = self._diff_label.center(len(label_cells[:-1]), ' ')
189        label_rows.extend(
190            [
191                f"{self._cs.TL.value}{divider_cells[:-1]}{self._cs.TR.value}",
192                f"{self._cs.V.value}{label_cells}{self._cs.V.value}",
193                f"{self._cs.ML.value}{divider_cells[:-1]}{self._cs.MR.value}",
194            ]
195        )
196        return label_rows
197
198    def create_table(self) -> str:
199        """Parse DataSourceMap to create ASCII table."""
200        curr_lb_hierarchy = None
201        last_diff_name = ''
202        if self._diff_mode:
203            self._ascii_table_rows.extend([*self._label_title_row()])
204        else:
205            self._ascii_table_rows.extend(
206                [self._create_border(True, self._cs.H.value)]
207            )
208        self._ascii_table_rows.extend([*self._create_title_row()])
209
210        has_entries = False
211
212        for curr_label in self._data_source_map.labels():
213            if curr_label.size == 0:
214                continue
215
216            has_entries = True
217
218            new_lb_hierarchy = tuple(
219                [
220                    *self._get_ds_label_size(curr_label.parents),
221                    self._LabelContent(
222                        curr_label.name,
223                        curr_label.size,
224                        get_label_status(curr_label),
225                    ),
226                ]
227            )
228            diff_list = self._diff_label_names(
229                curr_lb_hierarchy, new_lb_hierarchy
230            )
231            curr_lb_hierarchy = new_lb_hierarchy
232
233            if curr_label.parents and curr_label.parents[0] == last_diff_name:
234                continue
235            if (
236                self._diff_mode
237                and diff_list[0].name
238                and (
239                    not cast(
240                        DiffDataSourceMap, self._data_source_map
241                    ).has_diff_sublabels(diff_list[0].name)
242                )
243            ):
244                if (len(self._ascii_table_rows) > 5) and (
245                    self._ascii_table_rows[-1][0] != '+'
246                ):
247                    self._ascii_table_rows.append(
248                        self._row_divider(
249                            len(self._col_names), self._cs.H.value
250                        )
251                    )
252                self._ascii_table_rows.append(
253                    self._create_same_label_row(1, diff_list[0].name)
254                )
255
256                last_diff_name = curr_label.parents[0]
257            else:
258                self._ascii_table_rows += self._create_diff_rows(diff_list)
259
260        if self._rst_output and self._ascii_table_rows[-1][0] == '+':
261            self._ascii_table_rows.pop()
262
263        self._ascii_table_rows.extend(
264            [*self._create_total_row(is_empty=not has_entries)]
265        )
266
267        return '\n'.join(self._ascii_table_rows) + '\n'
268
269    def _create_same_label_row(self, col_index: int, label: str) -> str:
270        label_row = ''
271        for col in range(len(self._col_names) - 1):
272            if col == col_index:
273                curr_cell = self._create_cell(label, False, col, _Align.LEFT)
274            else:
275                curr_cell = self._create_cell('', False, col)
276            label_row += curr_cell
277        label_row += self._create_cell(
278            "(SAME)", True, len(self._col_widths) - 1, _Align.RIGHT
279        )
280        return label_row
281
282    def _get_ds_label_size(
283        self, parent_labels: tuple[str, ...]
284    ) -> Iterable[_LabelContent]:
285        """Produce label, size pairs from parent label names."""
286        parent_label_sizes = []
287        for index, target_label in enumerate(parent_labels):
288            for curr_label in self._data_source_map.labels(index):
289                if curr_label.name == target_label:
290                    diff_label = get_label_status(curr_label)
291                    parent_label_sizes.append(
292                        self._LabelContent(
293                            curr_label.name, curr_label.size, diff_label
294                        )
295                    )
296                    break
297        return parent_label_sizes
298
299    def _create_total_row(self, is_empty: bool = False) -> Iterable[str]:
300        complete_total_rows = []
301
302        if self._diff_mode and is_empty:
303            # When diffing two identical binaries, output a row indicating that
304            # the two are the same.
305            no_diff_row = ''
306            for i in range(len(self._col_names)):
307                if i == 0:
308                    no_diff_row += self._create_cell(
309                        'N/A', False, i, _Align.CENTER
310                    )
311                elif i == len(self._col_names) - 1:
312                    no_diff_row += self._create_cell('0', True, i)
313                else:
314                    no_diff_row += self._create_cell(
315                        '(same)', False, i, _Align.CENTER
316                    )
317            complete_total_rows.append(no_diff_row)
318
319        complete_total_rows.append(
320            self._row_divider(len(self._col_names), self._total_divider)
321        )
322        total_row = ''
323
324        for i in range(len(self._col_names)):
325            if i == 0:
326                total_row += self._create_cell('Total', False, i, _Align.LEFT)
327            elif i == len(self._col_names) - 1:
328                total_size_str = diff_sign_sizes(
329                    self._total_size, self._diff_mode
330                )
331                total_row += self._create_cell(total_size_str, True, i)
332            else:
333                total_row += self._create_cell('', False, i, _Align.CENTER)
334
335        complete_total_rows.extend(
336            [total_row, self._create_border(False, self._cs.H.value)]
337        )
338        return complete_total_rows
339
340    def _create_diff_rows(
341        self, diff_list: tuple[_LabelContent, ...]
342    ) -> Iterable[str]:
343        """Create rows for each label according to its index in diff_list."""
344        curr_row = ''
345        diff_index = 0
346        diff_rows = []
347        for index, label_content in enumerate(diff_list):
348            if label_content.name:
349                if self._diff_mode:
350                    curr_row += self._create_cell(
351                        label_content.label_status, False, 0
352                    )
353                    diff_index = 1
354                for cell_index in range(
355                    diff_index, len(diff_list) + diff_index
356                ):
357                    if cell_index == index + diff_index:
358                        if (
359                            cell_index == diff_index
360                            and len(self._ascii_table_rows) > 5
361                            and not self._rst_output
362                        ):
363                            diff_rows.append(
364                                self._row_divider(
365                                    len(self._col_names), self._cs.H.value
366                                )
367                            )
368                        if (
369                            len(label_content.name) + self._additional_padding
370                        ) > self._col_widths[cell_index]:
371                            curr_row = self._multi_row_label(
372                                label_content.name, cell_index
373                            )
374                            break
375                        curr_row += self._create_cell(
376                            label_content.name, False, cell_index, _Align.LEFT
377                        )
378                    else:
379                        curr_row += self._create_cell('', False, cell_index)
380
381                # Add size end of current row.
382                curr_size = diff_sign_sizes(label_content.size, self._diff_mode)
383                curr_row += self._create_cell(
384                    curr_size, True, len(self._col_widths) - 1, _Align.RIGHT
385                )
386                diff_rows.append(curr_row)
387                if self._rst_output:
388                    diff_rows.append(
389                        self._row_divider(
390                            len(self._col_names), self._cs.H.value
391                        )
392                    )
393                curr_row = ''
394
395        return diff_rows
396
397    def _create_cell(
398        self,
399        content: str,
400        last_cell: bool,
401        col_index: int,
402        align: _Align | None = _Align.RIGHT,
403    ) -> str:
404        v_border = self._cs.V.value
405        if self._rst_output and content:
406            content = content.replace('_', '\\_')
407        pad_diff = self._col_widths[col_index] - len(content)
408        padding = (pad_diff // 2) * ' '
409        odd_pad = ' ' if pad_diff % 2 == 1 else ''
410        string_cell = ''
411
412        if align == _Align.CENTER:
413            string_cell = f'{v_border}{odd_pad}{padding}{content}{padding}'
414        elif align == _Align.LEFT:
415            string_cell = f'{v_border}{content}{padding*2}{odd_pad}'
416        elif align == _Align.RIGHT:
417            string_cell = f'{v_border}{padding*2}{odd_pad}{content}'
418
419        if last_cell:
420            string_cell += self._cs.V.value
421        return string_cell
422
423    def _multi_row_label(self, content: str, target_col_index: int) -> str:
424        """Split content name into multiple rows within correct column."""
425        max_len = self._col_widths[target_col_index] - self._additional_padding
426        split_content = '...'.join(
427            content[max_len:][i : i + max_len - 3]
428            for i in range(0, len(content[max_len:]), max_len - 3)
429        )
430        split_content = f"{content[:max_len]}...{split_content}"
431        split_tab_content = [
432            split_content[i : i + max_len]
433            for i in range(0, len(split_content), max_len)
434        ]
435        multi_label = []
436        curr_row = ''
437        for index, cut_content in enumerate(split_tab_content):
438            last_cell = False
439            for blank_cell_index in range(len(self._col_names)):
440                if blank_cell_index == target_col_index:
441                    curr_row += self._create_cell(
442                        cut_content, False, target_col_index, _Align.LEFT
443                    )
444                else:
445                    if blank_cell_index == len(self._col_names) - 1:
446                        if index == len(split_tab_content) - 1:
447                            break
448                        last_cell = True
449                    curr_row += self._create_cell(
450                        '', last_cell, blank_cell_index
451                    )
452            multi_label.append(curr_row)
453            curr_row = ''
454
455        return '\n'.join(multi_label)
456
457    def _row_divider(self, col_num: int, h_div: str) -> str:
458        l_border = ''
459        r_border = ''
460        row_div = ''
461        for col in range(col_num):
462            if col == 0:
463                l_border = self._cs.ML.value
464                r_border = ''
465            elif col == (col_num - 1):
466                l_border = self._cs.MM.value
467                r_border = self._cs.MR.value
468            else:
469                l_border = self._cs.MM.value
470                r_border = ''
471
472            row_div += f"{l_border}{self._col_widths[col] * h_div}{r_border}"
473        return row_div
474
475    def _create_title_row(self) -> Iterable[str]:
476        title_rows = []
477        title_cells = ''
478        last_cell = False
479        for index, curr_name in enumerate(self._col_names):
480            if index == len(self._col_names) - 1:
481                last_cell = True
482            title_cells += self._create_cell(
483                curr_name, last_cell, index, _Align.CENTER
484            )
485        title_rows.extend(
486            [
487                title_cells,
488                self._row_divider(len(self._col_names), self._cs.HH.value),
489            ]
490        )
491        return title_rows
492
493    def _create_border(self, top: bool, h_div: str):
494        """Top or bottom borders of ASCII table."""
495        row_div = ''
496        for col in range(len(self._col_names)):
497            if top:
498                if col == 0:
499                    l_div = self._cs.TL.value
500                    r_div = ''
501                elif col == (len(self._col_names) - 1):
502                    l_div = self._cs.TM.value
503                    r_div = self._cs.TR.value
504                else:
505                    l_div = self._cs.TM.value
506                    r_div = ''
507            else:
508                if col == 0:
509                    l_div = self._cs.BL.value
510                    r_div = ''
511                elif col == (len(self._col_names) - 1):
512                    l_div = self._cs.BM.value
513                    r_div = self._cs.BR.value
514                else:
515                    l_div = self._cs.BM.value
516                    r_div = ''
517
518            row_div += f"{l_div}{self._col_widths[col] * h_div}{r_div}"
519        return row_div
520
521
522class RstOutput:
523    """Tabular output in ASCII format, which is also valid RST."""
524
525    def __init__(self, ds_map: DataSourceMap, table_label: str | None = None):
526        self._data_source_map = ds_map
527        self._table_label = table_label
528        self._diff_mode = False
529        if isinstance(self._data_source_map, DiffDataSourceMap):
530            self._diff_mode = True
531
532    def create_table(self) -> str:
533        """Initializes RST table and builds first row."""
534        table_builder = [
535            '\n.. list-table::',
536            '   :widths: auto',
537            '   :header-rows: 1\n',
538        ]
539        header_cols = ['Label', 'Segment', 'Delta']
540        for i, col_name in enumerate(header_cols):
541            list_space = '*' if i == 0 else ' '
542            table_builder.append(f"   {list_space} - {col_name}")
543
544        return '\n'.join(table_builder) + f'\n{self.add_report_row()}\n'
545
546    def _label_status_unchanged(self, parent_lb_name: str) -> bool:
547        """Determines if parent label has no status change in diff mode."""
548        for curr_lb in self._data_source_map.labels():
549            if curr_lb.size != 0:
550                if (
551                    curr_lb.parents and (parent_lb_name == curr_lb.parents[0])
552                ) or (curr_lb.name == parent_lb_name):
553                    if get_label_status(curr_lb) != '':
554                        return False
555        return True
556
557    def add_report_row(self) -> str:
558        """Add in new size report row with Label, Segment, and Delta.
559
560        Returns:
561            RST string that is the current row with a full symbols
562            table breakdown of the corresponding segment.
563        """
564        table_rows = []
565        curr_row = []
566        curr_label_name = ''
567        for parent_lb in self._data_source_map.labels(0):
568            if parent_lb.size != 0:
569                if (self._table_label is not None) and (
570                    curr_label_name != self._table_label
571                ):
572                    curr_row.append(f'   * - {self._table_label} ')
573                    curr_label_name = self._table_label
574                else:
575                    curr_row.append('   * -')
576                curr_row.extend(
577                    [
578                        f'     - .. dropdown:: {parent_lb.name}',
579                        '            :animate: fade-in\n',
580                        '            .. list-table::',
581                        '               :widths: auto\n',
582                    ]
583                )
584                if self._label_status_unchanged(parent_lb.name):
585                    skip_status = 1, '*'
586                else:
587                    skip_status = 0, ' '
588                for curr_lb in self._data_source_map.labels():
589                    if (curr_lb.size != 0) and (
590                        (
591                            curr_lb.parents
592                            and (parent_lb.name == curr_lb.parents[0])
593                        )
594                        or (curr_lb.name == parent_lb.name)
595                    ):
596                        sign_size = diff_sign_sizes(
597                            curr_lb.size, self._diff_mode
598                        )
599                        curr_status = get_label_status(curr_lb)
600                        curr_name = curr_lb.name.replace('_', '\\_')
601                        to_extend = [
602                            f'               * - {curr_status}',
603                            f'               {skip_status[1]} - {sign_size}',
604                            f'                 - {curr_name}\n',
605                        ][skip_status[0] :]
606                        curr_row.extend(to_extend)
607                curr_row.append(
608                    f'     - {diff_sign_sizes(parent_lb.size, self._diff_mode)}'
609                )
610            table_rows.extend(curr_row)
611            curr_row = []
612
613        # No size difference.
614        if len(table_rows) == 0:
615            table_rows.extend(
616                [f'\n   * - {self._table_label}', '     - (ALL)', '     - 0']
617            )
618
619        return '\n'.join(table_rows)
620