xref: /aosp_15_r20/external/angle/build/android/method_count.py (revision 8975f5c5ed3d1c378011245431ada316dfb6f244)
1#! /usr/bin/env python3
2# Copyright 2015 The Chromium Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6
7import argparse
8import os
9import re
10import zipfile
11
12from pylib.dex import dex_parser
13
14
15class DexStatsCollector:
16  """Tracks count of method/field/string/type as well as unique methods."""
17
18  def __init__(self):
19    # Signatures of all methods from all seen dex files.
20    self._unique_methods = set()
21    # Map of label -> { metric -> count }.
22    self._counts_by_label = {}
23
24  def _CollectFromDexfile(self, label, dexfile):
25    assert label not in self._counts_by_label, 'exists: ' + label
26    self._counts_by_label[label] = {
27        'fields': dexfile.header.field_ids_size,
28        'methods': dexfile.header.method_ids_size,
29        'strings': dexfile.header.string_ids_size,
30        'types': dexfile.header.type_ids_size,
31    }
32    self._unique_methods.update(dexfile.IterMethodSignatureParts())
33
34  def CollectFromZip(self, label, path):
35    """Add dex stats from an .apk/.jar/.aab/.zip."""
36    with zipfile.ZipFile(path, 'r') as z:
37      for subpath in z.namelist():
38        if not re.match(r'.*classes\d*\.dex$', subpath):
39          continue
40        dexfile = dex_parser.DexFile(bytearray(z.read(subpath)))
41        self._CollectFromDexfile('{}!{}'.format(label, subpath), dexfile)
42
43  def CollectFromDex(self, label, path):
44    """Add dex stats from a .dex file."""
45    with open(path, 'rb') as f:
46      dexfile = dex_parser.DexFile(bytearray(f.read()))
47    self._CollectFromDexfile(label, dexfile)
48
49  def MergeFrom(self, parent_label, other):
50    """Add dex stats from another DexStatsCollector."""
51    # pylint: disable=protected-access
52    for label, other_counts in other._counts_by_label.items():
53      new_label = '{}-{}'.format(parent_label, label)
54      self._counts_by_label[new_label] = other_counts.copy()
55    self._unique_methods.update(other._unique_methods)
56    # pylint: enable=protected-access
57
58  def GetUniqueMethodCount(self):
59    """Returns total number of unique methods across encountered dex files."""
60    return len(self._unique_methods)
61
62  def GetCountsByLabel(self):
63    """Returns dict of label -> {metric -> count}."""
64    return self._counts_by_label
65
66  def GetTotalCounts(self):
67    """Returns dict of {metric -> count}, where |count| is sum(metric)."""
68    ret = {}
69    for metric in ('fields', 'methods', 'strings', 'types'):
70      ret[metric] = sum(x[metric] for x in self._counts_by_label.values())
71    return ret
72
73  def GetDexCacheSize(self, pre_oreo):
74    """Returns number of bytes of dirty RAM is consumed from all dex files."""
75    # Dex Cache was optimized in Android Oreo:
76    # https://source.android.com/devices/tech/dalvik/improvements#dex-cache-removal
77    if pre_oreo:
78      total = sum(self.GetTotalCounts().values())
79    else:
80      total = sum(c['methods'] for c in self._counts_by_label.values())
81    return total * 4  # 4 bytes per entry.
82
83
84def main():
85  parser = argparse.ArgumentParser()
86  parser.add_argument('paths', nargs='+')
87  args = parser.parse_args()
88
89  collector = DexStatsCollector()
90  for path in args.paths:
91    if os.path.splitext(path)[1] in ('.zip', '.apk', '.jar', '.aab'):
92      collector.CollectFromZip(path, path)
93    else:
94      collector.CollectFromDex(path, path)
95
96  counts_by_label = collector.GetCountsByLabel()
97  for label, counts in sorted(counts_by_label.items()):
98    print('{}:'.format(label))
99    for metric, count in sorted(counts.items()):
100      print('  {}:'.format(metric), count)
101    print()
102
103  if len(counts_by_label) > 1:
104    print('Totals:')
105    for metric, count in sorted(collector.GetTotalCounts().items()):
106      print('  {}:'.format(metric), count)
107    print()
108
109  print('Unique Methods:', collector.GetUniqueMethodCount())
110  print('DexCache (Pre-Oreo):', collector.GetDexCacheSize(pre_oreo=True),
111        'bytes of dirty memory')
112  print('DexCache (Oreo+):', collector.GetDexCacheSize(pre_oreo=False),
113        'bytes of dirty memory')
114
115
116if __name__ == '__main__':
117  main()
118