#!/usr/bin/env vpython
# Copyright 2017 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.
"""This script helps to generate code coverage report.

  It uses Clang Source-based Code Coverage -
  https://clang.llvm.org/docs/SourceBasedCodeCoverage.html

  In order to generate code coverage report, you need to first add
  "use_clang_coverage=true" and "is_component_build=false" GN flags to args.gn
  file in your build output directory (e.g. out/coverage).

  * Example usage:

  gn gen out/coverage \\
      --args="use_clang_coverage=true is_component_build=false\\
              is_debug=false dcheck_always_on=true"
  gclient runhooks
  vpython tools/code_coverage/coverage.py crypto_unittests url_unittests \\
      -b out/coverage -o out/report -c 'out/coverage/crypto_unittests' \\
      -c 'out/coverage/url_unittests --gtest_filter=URLParser.PathURL' \\
      -f url/ -f crypto/

  The command above builds crypto_unittests and url_unittests targets and then
  runs them with specified command line arguments. For url_unittests, it only
  runs the test URLParser.PathURL. The coverage report is filtered to include
  only files and sub-directories under url/ and crypto/ directories.

  If you want to run tests that try to draw to the screen but don't have a
  display connected, you can run tests in headless mode with xvfb.

  * Sample flow for running a test target with xvfb (e.g. unit_tests):

  vpython tools/code_coverage/coverage.py unit_tests -b out/coverage \\
      -o out/report -c 'python testing/xvfb.py out/coverage/unit_tests'

  If you are building a fuzz target, you need to add "use_libfuzzer=true" GN
  flag as well.

  * Sample workflow for a fuzz target (e.g. pdfium_fuzzer):

  vpython tools/code_coverage/coverage.py pdfium_fuzzer \\
      -b out/coverage -o out/report \\
      -c 'out/coverage/pdfium_fuzzer -runs=0 <corpus_dir>' \\
      -f third_party/pdfium

  where:
    <corpus_dir> - directory containing samples files for this format.

  To learn more about generating code coverage reports for fuzz targets, see
  https://chromium.googlesource.com/chromium/src/+/main/testing/libfuzzer/efficient_fuzzer.md#Code-Coverage

  * Sample workflow for running Blink web tests:

  vpython tools/code_coverage/coverage.py blink_tests \\
      -wt -b out/coverage -o out/report -f third_party/blink

  If you need to pass arguments to run_web_tests.py, use
    -wt='arguments to run_web_tests.py e.g. test directories'

  For more options, please refer to tools/code_coverage/coverage.py -h.

  For an overview of how code coverage works in Chromium, please refer to
  https://chromium.googlesource.com/chromium/src/+/main/docs/testing/code_coverage.md
"""

from __future__ import print_function

import sys

import argparse
import json
import logging
import multiprocessing
import os
import platform
import re
import shlex
import shutil
import subprocess
import six

if six.PY2:
  from urllib2 import urlopen
else:
  from urllib.request import urlopen

sys.path.append(
    os.path.join(
        os.path.dirname(__file__), os.path.pardir, os.path.pardir,
        'third_party'))
from collections import defaultdict

import coverage_utils

# Absolute path to the code coverage tools binary. These paths can be
# overwritten by user specified coverage tool paths.
# Absolute path to the root of the checkout.
SRC_ROOT_PATH = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                             os.path.pardir, os.path.pardir)
LLVM_BIN_DIR = os.path.join(
    os.path.join(SRC_ROOT_PATH, 'third_party', 'llvm-build', 'Release+Asserts'),
    'bin')
LLVM_COV_PATH = os.path.join(LLVM_BIN_DIR, 'llvm-cov')
LLVM_PROFDATA_PATH = os.path.join(LLVM_BIN_DIR, 'llvm-profdata')


# Build directory, the value is parsed from command line arguments.
BUILD_DIR = None

# Output directory for generated artifacts, the value is parsed from command
# line arguemnts.
OUTPUT_DIR = None

# Name of the file extension for profraw data files.
PROFRAW_FILE_EXTENSION = 'profraw'

# Name of the final profdata file, and this file needs to be passed to
# "llvm-cov" command in order to call "llvm-cov show" to inspect the
# line-by-line coverage of specific files.
PROFDATA_FILE_NAME = os.extsep.join(['coverage', 'profdata'])

# Name of the file with summary information generated by llvm-cov export.
SUMMARY_FILE_NAME = os.extsep.join(['summary', 'json'])

# Name of the coverage file in lcov format generated by llvm-cov export.
LCOV_FILE_NAME = os.extsep.join(['coverage', 'lcov'])

# Build arg required for generating code coverage data.
CLANG_COVERAGE_BUILD_ARG = 'use_clang_coverage'

LOGS_DIR_NAME = 'logs'

# Used to extract a mapping between directories and components.
COMPONENT_MAPPING_URL = (
    'https://storage.googleapis.com/chromium-owners/component_map.json')

# Caches the results returned by _GetBuildArgs, don't use this variable
# directly, call _GetBuildArgs instead.
_BUILD_ARGS = None

# Retry failed merges.
MERGE_RETRIES = 3

# Message to guide user to file a bug when everything else fails.
FILE_BUG_MESSAGE = (
    'If it persists, please file a bug with the command you used, git revision '
    'and args.gn config here: '
    'https://bugs.chromium.org/p/chromium/issues/entry?'
    'components=Infra%3ETest%3ECodeCoverage')

# String to replace with actual llvm profile path.
LLVM_PROFILE_FILE_PATH_SUBSTITUTION = '<llvm_profile_file_path>'

def _ConfigureLLVMCoverageTools(args):
  """Configures llvm coverage tools."""
  if args.coverage_tools_dir:
    llvm_bin_dir = coverage_utils.GetFullPath(args.coverage_tools_dir)
    global LLVM_COV_PATH
    global LLVM_PROFDATA_PATH
    LLVM_COV_PATH = os.path.join(llvm_bin_dir, 'llvm-cov')
    LLVM_PROFDATA_PATH = os.path.join(llvm_bin_dir, 'llvm-profdata')
  else:
    subprocess.check_call([
        sys.executable, 'tools/clang/scripts/update.py', '--package',
        'coverage_tools'
    ])

  if coverage_utils.GetHostPlatform() == 'win':
    LLVM_COV_PATH += '.exe'
    LLVM_PROFDATA_PATH += '.exe'

  coverage_tools_exist = (
      os.path.exists(LLVM_COV_PATH) and os.path.exists(LLVM_PROFDATA_PATH))
  assert coverage_tools_exist, ('Cannot find coverage tools, please make sure '
                                'both \'%s\' and \'%s\' exist.') % (
                                    LLVM_COV_PATH, LLVM_PROFDATA_PATH)


def _GetPathWithLLVMSymbolizerDir():
  """Add llvm-symbolizer directory to path for symbolized stacks."""
  path = os.getenv('PATH')
  dirs = path.split(os.pathsep)
  if LLVM_BIN_DIR in dirs:
    return path

  return path + os.pathsep + LLVM_BIN_DIR


def _GetTargetOS():
  """Returns the target os specified in args.gn file.

  Returns an empty string is target_os is not specified.
  """
  build_args = _GetBuildArgs()
  return build_args['target_os'] if 'target_os' in build_args else ''


def _IsAndroid():
  """Returns true if the target_os specified in args.gn file is android"""
  return _GetTargetOS() == 'android'


def _IsIOS():
  """Returns true if the target_os specified in args.gn file is ios"""
  return _GetTargetOS() == 'ios'


def _GeneratePerFileLineByLineCoverageInFormat(binary_paths, profdata_file_path,
                                               filters, ignore_filename_regex,
                                               output_format):
  """Generates per file line-by-line coverage in html or text using
  'llvm-cov show'.

  For a file with absolute path /a/b/x.cc, a html/txt report is generated as:
  OUTPUT_DIR/coverage/a/b/x.cc.[html|txt]. For html format, an index html file
  is also generated as: OUTPUT_DIR/index.html.

  Args:
    binary_paths: A list of paths to the instrumented binaries.
    profdata_file_path: A path to the profdata file.
    filters: A list of directories and files to get coverage for.
    ignore_filename_regex: A regular expression for skipping source code files
                           with certain file paths.
    output_format: The output format of generated report files.
  """
  # llvm-cov show [options] -instr-profile PROFILE BIN [-object BIN,...]
  # [[-object BIN]] [SOURCES]
  # NOTE: For object files, the first one is specified as a positional argument,
  # and the rest are specified as keyword argument.
  logging.debug('Generating per file line by line coverage reports using '
                '"llvm-cov show" command.')

  subprocess_cmd = [
      LLVM_COV_PATH, 'show', '-format={}'.format(output_format),
      '-compilation-dir={}'.format(BUILD_DIR),
      '-output-dir={}'.format(OUTPUT_DIR),
      '-instr-profile={}'.format(profdata_file_path), binary_paths[0]
  ]
  subprocess_cmd.extend(
      ['-object=' + binary_path for binary_path in binary_paths[1:]])
  _AddArchArgumentForIOSIfNeeded(subprocess_cmd, len(binary_paths))
  if coverage_utils.GetHostPlatform() in ['linux', 'mac']:
    subprocess_cmd.extend(['-Xdemangler', 'c++filt', '-Xdemangler', '-n'])
  subprocess_cmd.extend(filters)
  if ignore_filename_regex:
    subprocess_cmd.append('-ignore-filename-regex=%s' % ignore_filename_regex)

  subprocess.check_call(subprocess_cmd)

  logging.debug('Finished running "llvm-cov show" command.')


def _GeneratePerFileLineByLineCoverageInLcov(binary_paths, profdata_file_path, filters,
                                             ignore_filename_regex):
  """Generates per file line-by-line coverage using "llvm-cov export".

  Args:
    binary_paths: A list of paths to the instrumented binaries.
    profdata_file_path: A path to the profdata file.
    filters: A list of directories and files to get coverage for.
    ignore_filename_regex: A regular expression for skipping source code files
                           with certain file paths.
  """
  logging.debug('Generating per file line by line coverage reports using '
                '"llvm-cov export" command.')
  for path in binary_paths:
    if not os.path.exists(path):
      logging.error("Binary %s does not exist", path)
  subprocess_cmd = [
      LLVM_COV_PATH, 'export', '-format=lcov',
      '-instr-profile=' + profdata_file_path, binary_paths[0]
  ]
  subprocess_cmd.extend(
      ['-object=' + binary_path for binary_path in binary_paths[1:]])
  _AddArchArgumentForIOSIfNeeded(subprocess_cmd, len(binary_paths))
  subprocess_cmd.extend(filters)
  if ignore_filename_regex:
    subprocess_cmd.append('-ignore-filename-regex=%s' % ignore_filename_regex)

  # Write output on the disk to be used by code coverage bot.
  with open(_GetLcovFilePath(), 'w') as f:
    subprocess.check_call(subprocess_cmd, stdout=f)

  logging.debug('Finished running "llvm-cov export" command.')


def _GetLogsDirectoryPath():
  """Path to the logs directory."""
  return os.path.join(
      coverage_utils.GetCoverageReportRootDirPath(OUTPUT_DIR), LOGS_DIR_NAME)


def _GetProfdataFilePath():
  """Path to the resulting .profdata file."""
  return os.path.join(
      coverage_utils.GetCoverageReportRootDirPath(OUTPUT_DIR),
      PROFDATA_FILE_NAME)


def _GetSummaryFilePath():
  """The JSON file that contains coverage summary written by llvm-cov export."""
  return os.path.join(
      coverage_utils.GetCoverageReportRootDirPath(OUTPUT_DIR),
      SUMMARY_FILE_NAME)


def _GetLcovFilePath():
  """The LCOV file that contains coverage data written by llvm-cov export."""
  return os.path.join(
      coverage_utils.GetCoverageReportRootDirPath(OUTPUT_DIR),
      LCOV_FILE_NAME)


def _CreateCoverageProfileDataForTargets(targets, commands, jobs_count=None):
  """Builds and runs target to generate the coverage profile data.

  Args:
    targets: A list of targets to build with coverage instrumentation.
    commands: A list of commands used to run the targets.
    jobs_count: Number of jobs to run in parallel for building. If None, a
                default value is derived based on CPUs availability.

  Returns:
    A relative path to the generated profdata file.
  """
  _BuildTargets(targets, jobs_count)
  target_profdata_file_paths = _GetTargetProfDataPathsByExecutingCommands(
      targets, commands)
  coverage_profdata_file_path = (
      _CreateCoverageProfileDataFromTargetProfDataFiles(
          target_profdata_file_paths))

  for target_profdata_file_path in target_profdata_file_paths:
    os.remove(target_profdata_file_path)

  return coverage_profdata_file_path


def _BuildTargets(targets, jobs_count):
  """Builds target with Clang coverage instrumentation.

  This function requires current working directory to be the root of checkout.

  Args:
    targets: A list of targets to build with coverage instrumentation.
    jobs_count: Number of jobs to run in parallel for compilation. If None, a
                default value is derived based on CPUs availability.
  """
  logging.info('Building %s.', str(targets))
  autoninja = 'autoninja'
  if coverage_utils.GetHostPlatform() == 'win':
    autoninja += '.bat'

  subprocess_cmd = [autoninja, '-C', BUILD_DIR]
  if jobs_count is not None:
    subprocess_cmd.append('-j' + str(jobs_count))

  subprocess_cmd.extend(targets)
  subprocess.check_call(subprocess_cmd)
  logging.debug('Finished building %s.', str(targets))


def _GetTargetProfDataPathsByExecutingCommands(targets, commands):
  """Runs commands and returns the relative paths to the profraw data files.

  Args:
    targets: A list of targets built with coverage instrumentation.
    commands: A list of commands used to run the targets.

  Returns:
    A list of relative paths to the generated profraw data files.
  """
  logging.debug('Executing the test commands.')

  # Remove existing profraw data files.
  report_root_dir = coverage_utils.GetCoverageReportRootDirPath(OUTPUT_DIR)
  for file_or_dir in os.listdir(report_root_dir):
    if file_or_dir.endswith(PROFRAW_FILE_EXTENSION):
      os.remove(os.path.join(report_root_dir, file_or_dir))

  # Ensure that logs directory exists.
  if not os.path.exists(_GetLogsDirectoryPath()):
    os.makedirs(_GetLogsDirectoryPath())

  profdata_file_paths = []

  # Run all test targets to generate profraw data files.
  for target, command in zip(targets, commands):
    output_file_name = os.extsep.join([target + '_output', 'log'])
    output_file_path = os.path.join(_GetLogsDirectoryPath(), output_file_name)

    profdata_file_path = None
    for _ in range(MERGE_RETRIES):
      logging.info('Running command: "%s", the output is redirected to "%s".',
                   command, output_file_path)

      if _IsIOSCommand(command):
        # On iOS platform, due to lack of write permissions, profraw files are
        # generated outside of the OUTPUT_DIR, and the exact paths are contained
        # in the output of the command execution.
        output = _ExecuteIOSCommand(command, output_file_path)
      else:
        # On other platforms, profraw files are generated inside the OUTPUT_DIR.
        output = _ExecuteCommand(target, command, output_file_path)

      profraw_file_paths = []
      if _IsIOS():
        profraw_file_paths = [_GetProfrawDataFileByParsingOutput(output)]
      elif _IsAndroid():
        android_coverage_dir = os.path.join(BUILD_DIR, 'coverage')
        for r, _, files in os.walk(android_coverage_dir):
          for f in files:
            if f.endswith(PROFRAW_FILE_EXTENSION):
              profraw_file_paths.append(os.path.join(r, f))
      else:
        for file_or_dir in os.listdir(report_root_dir):
          if file_or_dir.endswith(PROFRAW_FILE_EXTENSION):
            profraw_file_paths.append(
                os.path.join(report_root_dir, file_or_dir))

      assert profraw_file_paths, (
          'Running target "%s" failed to generate any profraw data file, '
          'please make sure the binary exists, is properly instrumented and '
          'does not crash. %s' % (target, FILE_BUG_MESSAGE))

      assert isinstance(profraw_file_paths, list), (
          'Variable \'profraw_file_paths\' is expected to be of type \'list\', '
          'but it is a %s. %s' % (type(profraw_file_paths), FILE_BUG_MESSAGE))

      try:
        profdata_file_path = _CreateTargetProfDataFileFromProfRawFiles(
            target, profraw_file_paths)
        break
      except Exception:
        logging.info('Retrying...')
      finally:
        # Remove profraw files now so that they are not used in next iteration.
        for profraw_file_path in profraw_file_paths:
          os.remove(profraw_file_path)

    assert profdata_file_path, (
        'Failed to merge target "%s" profraw files after %d retries. %s' %
        (target, MERGE_RETRIES, FILE_BUG_MESSAGE))
    profdata_file_paths.append(profdata_file_path)

  logging.debug('Finished executing the test commands.')

  return profdata_file_paths


def _GetEnvironmentVars(profraw_file_path):
  """Return environment vars for subprocess, given a profraw file path."""
  env = os.environ.copy()
  env.update({
      'LLVM_PROFILE_FILE': profraw_file_path,
      'PATH': _GetPathWithLLVMSymbolizerDir()
  })
  return env


def _SplitCommand(command):
  """Split a command string into parts in a platform-specific way."""
  if coverage_utils.GetHostPlatform() == 'win':
    return command.split()
  return shlex.split(command)


def _ExecuteCommand(target, command, output_file_path):
  """Runs a single command and generates a profraw data file."""
  # Per Clang "Source-based Code Coverage" doc:
  #
  # "%p" expands out to the process ID. It's not used by this scripts due to:
  # 1) If a target program spawns too many processess, it may exhaust all disk
  #    space available. For example, unit_tests writes thousands of .profraw
  #    files each of size 1GB+.
  # 2) If a target binary uses shared libraries, coverage profile data for them
  #    will be missing, resulting in incomplete coverage reports.
  #
  # "%Nm" expands out to the instrumented binary's signature. When this pattern
  # is specified, the runtime creates a pool of N raw profiles which are used
  # for on-line profile merging. The runtime takes care of selecting a raw
  # profile from the pool, locking it, and updating it before the program exits.
  # N must be between 1 and 9. The merge pool specifier can only occur once per
  # filename pattern.
  #
  # "%1m" is used when tests run in single process, such as fuzz targets.
  #
  # For other cases, "%4m" is chosen as it creates some level of parallelism,
  # but it's not too big to consume too much computing resource or disk space.
  profile_pattern_string = '%1m' if _IsFuzzerTarget(target) else '%4m'
  expected_profraw_file_name = os.extsep.join(
      [target, profile_pattern_string, PROFRAW_FILE_EXTENSION])
  expected_profraw_file_path = os.path.join(
      coverage_utils.GetCoverageReportRootDirPath(OUTPUT_DIR),
      expected_profraw_file_name)
  command = command.replace(LLVM_PROFILE_FILE_PATH_SUBSTITUTION,
                            expected_profraw_file_path)

  try:
    # Some fuzz targets or tests may write into stderr, redirect it as well.
    with open(output_file_path, 'wb') as output_file_handle:
      subprocess.check_call(_SplitCommand(command),
                            stdout=output_file_handle,
                            stderr=subprocess.STDOUT,
                            env=_GetEnvironmentVars(expected_profraw_file_path))
  except subprocess.CalledProcessError as e:
    logging.warning('Command: "%s" exited with non-zero return code.', command)

  return open(output_file_path, 'rb').read()


def _IsFuzzerTarget(target):
  """Returns true if the target is a fuzzer target."""
  build_args = _GetBuildArgs()
  use_libfuzzer = ('use_libfuzzer' in build_args and
                   build_args['use_libfuzzer'] == 'true')
  return use_libfuzzer and target.endswith('_fuzzer')


def _ExecuteIOSCommand(command, output_file_path):
  """Runs a single iOS command and generates a profraw data file.

  iOS application doesn't have write access to folders outside of the app, so
  it's impossible to instruct the app to flush the profraw data file to the
  desired location. The profraw data file will be generated somewhere within the
  application's Documents folder, and the full path can be obtained by parsing
  the output.
  """
  assert _IsIOSCommand(command)

  # After running tests, iossim generates a profraw data file, it won't be
  # needed anyway, so dump it into the OUTPUT_DIR to avoid polluting the
  # checkout.
  iossim_profraw_file_path = os.path.join(
      OUTPUT_DIR, os.extsep.join(['iossim', PROFRAW_FILE_EXTENSION]))
  command = command.replace(LLVM_PROFILE_FILE_PATH_SUBSTITUTION,
                            iossim_profraw_file_path)

  try:
    with open(output_file_path, 'wb') as output_file_handle:
      subprocess.check_call(_SplitCommand(command),
                            stdout=output_file_handle,
                            stderr=subprocess.STDOUT,
                            env=_GetEnvironmentVars(iossim_profraw_file_path))
  except subprocess.CalledProcessError as e:
    # iossim emits non-zero return code even if tests run successfully, so
    # ignore the return code.
    pass

  return open(output_file_path, 'rb').read()


def _GetProfrawDataFileByParsingOutput(output):
  """Returns the path to the profraw data file obtained by parsing the output.

  The output of running the test target has no format, but it is guaranteed to
  have a single line containing the path to the generated profraw data file.
  NOTE: This should only be called when target os is iOS.