# Copyright 2019 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Functions for interacting with llvm-profdata"""

import logging
import multiprocessing
import os
import subprocess

logging.basicConfig(
    format='[%(asctime)s %(levelname)s] %(message)s', level=logging.DEBUG)


def _call_profdata_tool(profile_input_file_paths,
                        profile_output_file_path,
                        profdata_tool_path,
                        retries=3):
  """Calls the llvm-profdata tool.

  Args:
    profile_input_file_paths: A list of relative paths to the files that
        are to be merged.
    profile_output_file_path: The path to the merged file to write.
    profdata_tool_path: The path to the llvm-profdata executable.

  Returns:
    A list of paths to profiles that had to be excluded to get the merge to
    succeed, suspected of being corrupted or malformed.

  Raises:
    CalledProcessError: An error occurred merging profiles.
  """
  logging.info('Merging profiles.')

  try:
    subprocess_cmd = [
        profdata_tool_path, 'merge', '-o', profile_output_file_path,
        '-sparse=true'
    ]
    subprocess_cmd.extend(profile_input_file_paths)

    # Redirecting stderr is required because when error happens, llvm-profdata
    # writes the error output to stderr and our error handling logic relies on
    # that output.
    output = subprocess.check_output(subprocess_cmd, stderr=subprocess.STDOUT)
    logging.info('Merge succeeded with output: %r', output)
  except subprocess.CalledProcessError as error:
    if len(profile_input_file_paths) > 1 and retries >= 0:
      logging.warning('Merge failed with error output: %r', error.output)

      # The output of the llvm-profdata command will include the path of
      # malformed files, such as
      # `error: /.../default.profraw: Malformed instrumentation profile data`
      invalid_profiles = [
          f for f in profile_input_file_paths if f in error.output
      ]

      if not invalid_profiles:
        logging.info(
            'Merge failed, but wasn\'t able to figure out the culprit invalid '
            'profiles from the output, so skip retry and bail out.')
        raise error

      valid_profiles = list(
          set(profile_input_file_paths) - set(invalid_profiles))
      if valid_profiles:
        logging.warning(
            'Following invalid profiles are removed as they were mentioned in '
            'the merge error output: %r', invalid_profiles)
        logging.info('Retry merging with the remaining profiles: %r',
                     valid_profiles)
        return invalid_profiles + _call_profdata_tool(
            valid_profiles, profile_output_file_path, profdata_tool_path,
            retries - 1)

    logging.error('Failed to merge profiles, return code (%d), output: %r' %
                  (error.returncode, error.output))
    raise error

  logging.info('Profile data is created as: "%r".', profile_output_file_path)
  return []


def _get_profile_paths(input_dir, input_extension):
  """Finds all the profiles in the given directory (recursively)."""
  paths = []
  for dir_path, _sub_dirs, file_names in os.walk(input_dir):
    paths.extend([
        os.path.join(dir_path, fn)
        for fn in file_names
        if fn.endswith(input_extension)
    ])
  return paths


def _validate_and_convert_profraws(profraw_files, profdata_tool_path):
  """Validates and converts profraws to profdatas.

  For each given .profraw file in the input, this method first validates it by
  trying to convert it to an indexed .profdata file, and if the validation and
  conversion succeeds, the generated .profdata file will be included in the
  output, otherwise, won't.

  This method is mainly used to filter out invalid profraw files.

  Args:
    profraw_files: A list of .profraw paths.
    profdata_tool_path: The path to the llvm-profdata executable.

  Returns:
    A tulple:
      A list of converted .profdata files of *valid* profraw files.
      A list of *invalid* profraw files.
  """
  logging.info('Validating and converting .profraw files.')

  for profraw_file in profraw_files:
    if not profraw_file.endswith('.profraw'):
      raise RuntimeError('%r is expected to be a .profraw file.' % profraw_file)

  cpu_count = multiprocessing.cpu_count()
  counts = max(10, cpu_count - 5)  # Use 10+ processes, but leave 5 cpu cores.
  pool = multiprocessing.Pool(counts)
  output_profdata_files = multiprocessing.Manager().list()
  invalid_profraw_files = multiprocessing.Manager().list()

  for profraw_file in profraw_files:
    pool.apply_async(_validate_and_convert_profraw,
                     (profraw_file, output_profdata_files,
                      invalid_profraw_files, profdata_tool_path))

  pool.close()
  pool.join()

  # Remove inputs, as they won't be needed and they can be pretty large.
  for input_file in profraw_files:
    os.remove(input_file)

  return list(output_profdata_files), list(invalid_profraw_files)


def _validate_and_convert_profraw(profraw_file, output_profdata_files,
                                  invalid_profraw_files, profdata_tool_path):
  output_profdata_file = profraw_file.replace('.profraw', '.profdata')
  subprocess_cmd = [
      profdata_tool_path, 'merge', '-o', output_profdata_file, '-sparse=true',
      profraw_file
  ]

  try:
    # Redirecting stderr is required because when error happens, llvm-profdata
    # writes the error output to stderr and our error handling logic relies on
    # that output.
    output = subprocess.check_output(subprocess_cmd, stderr=subprocess.STDOUT)
    logging.info('Validating and converting %r to %r succeeded with output: %r',
                 profraw_file, output_profdata_file, output)
    output_profdata_files.append(output_profdata_file)
  except subprocess.CalledProcessError as error:
    logging.warning('Validating and converting %r to %r failed with output: %r',
                    profraw_file, output_profdata_file, error.output)
    invalid_profraw_files.append(profraw_file)


def merge_profiles(input_dir, output_file, input_extension, profdata_tool_path):
  """Merges the profiles produced by the shards using llvm-profdata.

  Args:
    input_dir (str): The path to traverse to find input profiles.
    output_file (str): Where to write the merged profile.
    input_extension (str): File extension to look for in the input_dir.
        e.g. '.profdata' or '.profraw'
    profdata_tool_path: The path to the llvm-profdata executable.
  Returns:
    The list of profiles that had to be excluded to get the merge to
    succeed.
  """
  profile_input_file_paths = _get_profile_paths(input_dir, input_extension)
  invalid_profraw_files = []
  if input_extension == '.profraw':
    profile_input_file_paths, invalid_profraw_files = (
        _validate_and_convert_profraws(profile_input_file_paths,
                                       profdata_tool_path))
    logging.info('List of converted .profdata files: %r',
                 profile_input_file_paths)
    logging.info((
        'List of invalid .profraw files that failed to validate and convert: %r'
    ), invalid_profraw_files)

  invalid_profdata_files = _call_profdata_tool(
      profile_input_file_paths=profile_input_file_paths,
      profile_output_file_path=output_file,
      profdata_tool_path=profdata_tool_path)

  # Remove inputs, as they won't be needed and they can be pretty large.
  for input_file in profile_input_file_paths:
    os.remove(input_file)

  return invalid_profraw_files + invalid_profdata_files