#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2019 The ChromiumOS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Script to do post-process on orderfile generated by C3.

The goal of this script is to take in an orderfile generated by C3, and do
the following post process:

1. Take in the result of nm command on Chrome binary to find out all the
Builtin functions and put them after the input symbols.

2. Put special markers "chrome_begin_ordered_code" and "chrome_end_ordered_code"
in begin and end of the file.

The results of the file is intended to be uploaded and consumed when linking
Chrome in ChromeOS.
"""

from __future__ import division, print_function

import argparse
import os
import sys


def _parse_nm_output(stream):
  for line in (line.rstrip() for line in stream):
    if not line:
      continue

    pieces = line.split()
    if len(pieces) != 3:
      continue

    _, ty, symbol = pieces
    if ty not in 'tT':
      continue

    # We'll sometimes see synthesized symbols that start with $. There isn't
    # much we can do about or with them, regrettably.
    if symbol.startswith('$'):
      continue

    yield symbol


def _remove_duplicates(iterable):
  seen = set()
  for item in iterable:
    if item in seen:
      continue
    seen.add(item)
    yield item


def run(c3_ordered_stream, chrome_nm_stream, output_stream):
  head_marker = 'chrome_begin_ordered_code'
  tail_marker = 'chrome_end_ordered_code'

  c3_ordered_syms = [x.strip() for x in c3_ordered_stream.readlines()]
  all_chrome_syms = set(_parse_nm_output(chrome_nm_stream))
  # Sort by name, so it's predictable. Otherwise, these should all land in the
  # same hugepage anyway, so order doesn't matter as much.
  builtin_syms = sorted(s for s in all_chrome_syms if s.startswith('Builtins_'))
  output = _remove_duplicates([head_marker] + c3_ordered_syms + builtin_syms +
                              [tail_marker])
  output_stream.write('\n'.join(output))


def main(argv):
  parser = argparse.ArgumentParser()
  parser.add_argument('--chrome_nm', required=True, dest='chrome_nm')
  parser.add_argument('--input', required=True, dest='input_file')
  parser.add_argument('--output', required=True, dest='output_file')

  options = parser.parse_args(argv)

  if not os.path.exists(options.input_file):
    sys.exit("Input orderfile doesn\'t exist.")

  with open(options.input_file) as in_stream, \
  open(options.chrome_nm) as chrome_nm_stream, \
  open(options.output_file, 'w') as out_stream:
    run(in_stream, chrome_nm_stream, out_stream)


if __name__ == '__main__':
  sys.exit(main(sys.argv[1:]))