xref: /aosp_15_r20/external/pigweed/pw_allocator/py/pw_allocator/heap_viewer.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2020 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"""Heap visualizer of ASCII characters.
15
16TODO: b/328648868 - This tool is in need of an update, and should be considered
17**experimental** until then.
18"""
19
20from __future__ import annotations
21
22import argparse
23import sys
24import math
25import logging
26from dataclasses import dataclass, field
27import coloredlogs  # type: ignore
28
29
30@dataclass
31class HeapBlock:
32    """Building blocks for memory chunk allocated at heap."""
33
34    size: int
35    mem_offset: int
36    next: HeapBlock | None = None
37
38
39@dataclass
40class HeapUsage:
41    """Contains a linked list of allocated HeapBlocks."""
42
43    # Use a default_factory to avoid mutable default value. See
44    # https://docs.python.org/3/library/dataclasses.html#mutable-default-values
45    begin: HeapBlock = field(default_factory=lambda: HeapBlock(0, 0))
46
47    def add_block(self, block):
48        cur_block = self.begin.next
49        prev_block = self.begin
50        while cur_block is not None:
51            if cur_block.mem_offset == block.mem_offset:
52                return
53            if cur_block.mem_offset < block.mem_offset:
54                prev_block = cur_block
55                cur_block = cur_block.next
56            else:
57                block.next = cur_block
58                prev_block.next = block
59                return
60        prev_block.next = block
61
62    def remove_block(self, address):
63        cur_block = self.begin.next
64        prev_block = self.begin
65        while cur_block is not None:
66            if cur_block.mem_offset == address:
67                prev_block.next = cur_block.next
68                return
69            if cur_block.mem_offset < address:
70                prev_block = cur_block
71                cur_block = cur_block.next
72            else:
73                return
74
75
76def add_parser_arguments(parser):
77    """Parse args."""
78    parser.add_argument(
79        '--dump-file',
80        help=(
81            'dump file that contains a list of malloc and '
82            'free instructions. The format should be as '
83            'follows: "m <size> <address>" on a line for '
84            'each malloc called and "f <address>" on a line '
85            'for each free called.'
86        ),
87        required=True,
88    )
89
90    parser.add_argument(
91        '--heap-low-address',
92        help=('lower address of the heap.'),
93        type=lambda x: int(x, 0),
94        required=True,
95    )
96
97    parser.add_argument(
98        '--heap-high-address',
99        help=('higher address of the heap.'),
100        type=lambda x: int(x, 0),
101        required=True,
102    )
103
104    parser.add_argument(
105        '--poison-enabled',
106        help=('if heap poison is enabled or not.'),
107        default=False,
108        action='store_true',
109    )
110
111    parser.add_argument(
112        '--pointer-size',
113        help=('size of pointer on the machine.'),
114        default=4,
115        type=lambda x: int(x, 0),
116    )
117
118
119_LEFT_HEADER_CHAR = '['
120_RIGHT_HEADER_CHAR = ']'
121_USED_CHAR = '*'
122_FREE_CHAR = ' '
123_CHARACTERS_PER_LINE = 64
124_BYTES_PER_CHARACTER = 4
125_LOG = logging.getLogger(__name__)
126
127
128def _exit_due_to_file_not_found():
129    _LOG.critical(
130        'Dump file location is not provided or dump file is not '
131        'found. Please specify a valid file in the argument.'
132    )
133    sys.exit(1)
134
135
136def _exit_due_to_bad_heap_info():
137    _LOG.critical(
138        'Heap low/high address is missing or invalid. Please put valid '
139        'addresses in the argument.'
140    )
141    sys.exit(1)
142
143
144def visualize(
145    dump_file=None,
146    heap_low_address=None,
147    heap_high_address=None,
148    poison_enabled=False,
149    pointer_size=4,
150):
151    """Visualization of heap usage."""
152    # TODO: b/235282507 - Add standarized mechanisms to produce dump file and
153    # read heap information from dump file.
154    aligned_bytes = pointer_size
155    header_size = pointer_size * 2
156
157    try:
158        if heap_high_address < heap_low_address:
159            _exit_due_to_bad_heap_info()
160        heap_size = heap_high_address - heap_low_address
161    except TypeError:
162        _exit_due_to_bad_heap_info()
163
164    if poison_enabled:
165        poison_offset = pointer_size
166    else:
167        poison_offset = 0
168
169    try:
170        allocation_dump = open(dump_file, 'r')
171    except (FileNotFoundError, TypeError):
172        _exit_due_to_file_not_found()
173
174    heap_visualizer = HeapUsage()
175    # Parse the dump file.
176    for line in allocation_dump:
177        info = line[:-1].split(' ')
178        if info[0] == 'm':
179            # Add a HeapBlock when malloc is called
180            block = HeapBlock(
181                int(math.ceil(float(info[1]) / aligned_bytes)) * aligned_bytes,
182                int(info[2], 0) - heap_low_address,
183            )
184            heap_visualizer.add_block(block)
185        elif info[0] == 'f':
186            # Remove the HeapBlock when free is called
187            heap_visualizer.remove_block(int(info[1], 0) - heap_low_address)
188
189    # next_block indicates the nearest HeapBlock that hasn't finished
190    # printing.
191    next_block = heap_visualizer.begin
192    if next_block.next is None:
193        next_mem_offset = heap_size + header_size + poison_offset + 1
194        next_size = 0
195    else:
196        next_mem_offset = next_block.next.mem_offset
197        next_size = next_block.next.size
198
199    # Flags to indicate status of the 4 bytes going to be printed.
200    is_left_header = False
201    is_right_header = False
202    is_used = False
203
204    # Print overall heap information
205    _LOG.info(
206        '%-40s%-40s',
207        f'The heap starts at {hex(heap_low_address)}.',
208        f'The heap ends at {hex(heap_high_address)}.',
209    )
210    _LOG.info(
211        '%-40s%-40s',
212        f'Heap size is {heap_size // 1024}k bytes.',
213        f'Heap is aligned by {aligned_bytes} bytes.',
214    )
215    if poison_offset != 0:
216        _LOG.info(
217            'Poison is enabled %d bytes before and after the usable '
218            'space of each block.',
219            poison_offset,
220        )
221    else:
222        _LOG.info('%-40s', 'Poison is disabled.')
223    _LOG.info(
224        '%-40s',
225        'Below is the visualization of the heap. '
226        'Each character represents 4 bytes.',
227    )
228    _LOG.info('%-40s', f"    '{_FREE_CHAR}' indicates free space.")
229    _LOG.info('%-40s', f"    '{_USED_CHAR}' indicates used space.")
230    _LOG.info(
231        '%-40s',
232        f"    '{_LEFT_HEADER_CHAR}' indicates header or "
233        'poisoned space before the block.',
234    )
235    _LOG.info(
236        '%-40s',
237        f"    '{_RIGHT_HEADER_CHAR}' poisoned space after " 'the block.',
238    )
239    print()
240
241    # Go over the heap space where there will be 64 characters each line.
242    for line_base_address in range(
243        0, heap_size, _CHARACTERS_PER_LINE * _BYTES_PER_CHARACTER
244    ):
245        # Print the heap address of the current line.
246        sys.stdout.write(
247            f"{' ': <13}"
248            f'{hex(heap_low_address + line_base_address)}'
249            f"{f' (+{line_base_address}):': <12}"
250        )
251        for line_offset in range(
252            0, _CHARACTERS_PER_LINE * _BYTES_PER_CHARACTER, _BYTES_PER_CHARACTER
253        ):
254            # Determine if the current 4 bytes is used, unused, or is a
255            # header.
256            # The case that we have went over the previous block and will
257            # turn to the next block.
258            current_address = line_base_address + line_offset
259            if current_address == next_mem_offset + next_size + poison_offset:
260                next_block = next_block.next
261                # If this is the last block, set nextMemOffset to be over
262                # the last byte of heap so that the rest of the heap will
263                # be printed out as unused.
264                # Otherwise set the next HeapBlock allocated.
265                if next_block.next is None:
266                    next_mem_offset = (
267                        heap_size + header_size + poison_offset + 1
268                    )
269                    next_size = 0
270                else:
271                    next_mem_offset = next_block.next.mem_offset
272                    next_size = next_block.next.size
273
274            # Determine the status of the current 4 bytes.
275            if (
276                next_mem_offset - header_size - poison_offset
277                <= current_address
278                < next_mem_offset
279            ):
280                is_left_header = True
281                is_right_header = False
282                is_used = False
283            elif (
284                next_mem_offset <= current_address < next_mem_offset + next_size
285            ):
286                is_left_header = False
287                is_right_header = False
288                is_used = True
289            elif (
290                next_mem_offset + next_size
291                <= current_address
292                < next_mem_offset + next_size + poison_offset
293            ):
294                is_left_header = False
295                is_right_header = True
296                is_used = False
297            else:
298                is_left_header = False
299                is_right_header = False
300                is_used = False
301
302            if is_left_header:
303                sys.stdout.write(_LEFT_HEADER_CHAR)
304            elif is_right_header:
305                sys.stdout.write(_RIGHT_HEADER_CHAR)
306            elif is_used:
307                sys.stdout.write(_USED_CHAR)
308            else:
309                sys.stdout.write(_FREE_CHAR)
310        sys.stdout.write('\n')
311
312    allocation_dump.close()
313
314
315def main():
316    """A python script to visualize heap usage given a dump file."""
317    parser = argparse.ArgumentParser(description=main.__doc__)
318    add_parser_arguments(parser)
319    # Try to use pw_cli logs, else default to something reasonable.
320    try:
321        import pw_cli.log  # pylint: disable=import-outside-toplevel
322
323        pw_cli.log.install()
324    except ImportError:
325        coloredlogs.install(
326            level='INFO',
327            level_styles={'debug': {'color': 244}, 'error': {'color': 'red'}},
328            fmt='%(asctime)s %(levelname)s | %(message)s',
329        )
330    _LOG.warning("This tool is outdated and in need of an update!")
331    _LOG.warning("See b/328648868 for more details.")
332    visualize(**vars(parser.parse_args()))
333
334
335if __name__ == "__main__":
336    main()
337