Add TriChrome support to resource_size.py

Bug: 1059928
Change-Id: I8b39bc87abfecea5cc9f43a9cdd9cea3a228eed5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2108925
Commit-Queue: Andrew Grieve <agrieve@chromium.org>
Reviewed-by: Samuel Huang <huangs@chromium.org>
Cr-Commit-Position: refs/heads/master@{#751908}

build/android/resource_sizes.py

@@ -8,6 +8,8 @@
 More information at //docs/speed/binary_size/metrics.md.
 """
 
+from __future__ import print_function
+
 import argparse
 import collections
 from contextlib import contextmanager
@@ -577,27 +579,18 @@ def _ConfigOutDirAndToolsPrefix(out_dir):
   return out_dir, tool_prefix
 
 
-def _Analyze(apk_path, chartjson, args):
-
-  def report_func(*args):
-    # Do not add any new metrics without also documenting them in:
-    # //docs/speed/binary_size/metrics.md.
-    perf_tests_results_helper.ReportPerfResult(chartjson, *args)
-
+def _AnalyzeInternal(apk_path, report_func, args, apks_path=None):
   out_dir, tool_prefix = _ConfigOutDirAndToolsPrefix(args.out_dir)
-  apks_path = args.input if args.input.endswith('.apks') else None
   _DoApkAnalysis(apk_path, apks_path, tool_prefix, out_dir, report_func)
   _DoDexAnalysis(apk_path, report_func)
 
 
-def ResourceSizes(args):
-  chartjson = _BASE_CHART.copy() if args.output_format else None
-
-  if args.input.endswith('.apk'):
-    _Analyze(args.input, chartjson, args)
-  elif args.input.endswith('.apks'):
+def _AnalyzeApkOrApks(report_func, apk_path, args):
+  if apk_path.endswith('.apk'):
+    _AnalyzeInternal(apk_path, report_func, args)
+  elif apk_path.endswith('.apks'):
     with tempfile.NamedTemporaryFile(suffix='.apk') as f:
-      with zipfile.ZipFile(args.input) as z:
+      with zipfile.ZipFile(apk_path) as z:
         # Currently bundletool is creating two apks when .apks is created
         # without specifying an sdkVersion. Always measure the one with an
         # uncompressed shared library.
@@ -607,11 +600,55 @@ def ResourceSizes(args):
           info = z.getinfo('splits/base-master.apk')
         f.write(z.read(info))
         f.flush()
-      _Analyze(f.name, chartjson, args)
+      _AnalyzeInternal(f.name, report_func, args, apks_path=apk_path)
+  else:
+    raise Exception('Unknown file type: ' + apk_path)
+
+
+class _Reporter(object):
+  def __init__(self, chartjson):
+    self._chartjson = chartjson
+    self.trace_title_prefix = ''
+    self._combined_metrics = collections.defaultdict(int)
+
+  def __call__(self, graph_title, trace_title, value, units):
+    self._combined_metrics[(graph_title, trace_title, units)] += value
+
+    perf_tests_results_helper.ReportPerfResult(
+        self._chartjson, graph_title, self.trace_title_prefix + trace_title,
+        value, units)
+
+  def SynthesizeTotals(self):
+    for tup, value in sorted(self._combined_metrics.iteritems()):
+      graph_title, trace_title, units = tup
+      perf_tests_results_helper.ReportPerfResult(
+          self._chartjson, graph_title, 'Combined_' + trace_title, value, units)
+
+
+def _ResourceSizes(args):
+  chartjson = _BASE_CHART.copy() if args.output_format else None
+  reporter = _Reporter(chartjson)
+
+  specs = [
+      ('Chrome_', args.trichrome_chrome),
+      ('WebView_', args.trichrome_webview),
+      ('Library_', args.trichrome_library),
+  ]
+  for prefix, path in specs:
+    if path:
+      reporter.trace_title_prefix = prefix
+      _AnalyzeApkOrApks(reporter, path, args)
+
+  if any(path for _, path in specs):
+    reporter.SynthesizeTotals()
   else:
-    raise Exception('Unknown file type: ' + args.input)
+    _AnalyzeApkOrApks(reporter, args.input, args)
 
   if chartjson:
+    _DumpChartJson(args, chartjson)
+
+
+def _DumpChartJson(args, chartjson):
   if args.output_file == '-':
     json_file = sys.stdout
   elif args.output_file:
@@ -635,17 +672,14 @@ def ResourceSizes(args):
   if args.output_format == 'histograms':
     histogram_result = convert_chart_json.ConvertChartJson(results_path)
     if histogram_result.returncode != 0:
-        logging.error('chartjson conversion failed with error: %s',
+      raise Exception('chartjson conversion failed with error: ' +
                       histogram_result.stdout)
-        return 1
 
     histogram_path = os.path.join(args.output_dir, 'perf_results.json')
     logging.critical('Dumping histograms to %s', histogram_path)
     with open(histogram_path, 'w') as json_file:
       json_file.write(histogram_result.stdout)
 
-  return 0
-
 
 def main():
   argparser = argparse.ArgumentParser(description='Print APK size metrics.')
@@ -694,6 +728,15 @@ def main():
       'simplified JSON output format.')
 
   argparser.add_argument('input', help='Path to .apk or .apks file to measure.')
+  trichrome_group = argparser.add_argument_group(
+      'Trichrome inputs',
+      description='When specified, |input| is used only as Test suite name.')
+  trichrome_group.add_argument(
+      '--trichrome-chrome', help='Path to Trichrome Chrome .apks')
+  trichrome_group.add_argument(
+      '--trichrome-webview', help='Path to Trichrome WebView .apk(s)')
+  trichrome_group.add_argument(
+      '--trichrome-library', help='Path to Trichrome Library .apk')
   args = argparser.parse_args()
 
   devil_chromium.Initialize(output_directory=args.out_dir)
@@ -713,10 +756,10 @@ def main():
       os.makedirs(args.output_dir)
 
   try:
-    result = ResourceSizes(args)
+    _ResourceSizes(args)
     isolated_script_output = {
         'valid': True,
-        'failures': [test_name] if result else [],
+        'failures': [],
     }
   finally:
     if args.isolated_script_test_output:
@@ -726,8 +769,6 @@ def main():
       with open(args.isolated_script_test_output, 'w') as output_file:
         json.dump(isolated_script_output, output_file)
 
-  return result
-
 
 if __name__ == '__main__':
-  sys.exit(main())
+  main()