#!/usr/bin/env python3
# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Generates a PGO profile for LLVM.

**This script is meant to be run from inside of the chroot.**

Note that this script has a few (perhaps surprising) side-effects:
1. The first time this is run in a chroot, it will pack up your existing llvm
   and save it as a binpkg.
2. This script clobbers your llvm installation. If the script is run to
   completion, your old installation will be restored. If it does not, it may
   not be.
"""

import argparse
import dataclasses
import logging
import os
from pathlib import Path
import shlex
import shutil
import subprocess
import sys
import tempfile
import textwrap
from typing import Dict, FrozenSet, List, Optional

import pgo_tools


# This script runs `quickpkg` on LLVM. This file saves the version of LLVM that
# was quickpkg'ed.
SAVED_LLVM_BINPKG_STAMP = Path("/tmp/generate_pgo_profile_old_llvm.txt")

# Triple to build with when not trying to get backend coverage.
HOST_TRIPLE = "x86_64-pc-linux-gnu"

# List of triples we want coverage for.
IMPORTANT_TRIPLES = (
    HOST_TRIPLE,
    "x86_64-cros-linux-gnu",
    "armv7a-cros-linux-gnueabihf",
    "aarch64-cros-linux-gnu",
)

# Set of all of the cross-* libraries we need.
ALL_NEEDED_CROSS_LIBS = frozenset(
    f"cross-{triple}/{package}"
    for triple in IMPORTANT_TRIPLES
    if triple != HOST_TRIPLE
    for package in ("glibc", "libcxx", "llvm-libunwind", "linux-headers")
)


def ensure_llvm_binpkg_exists() -> bool:
    """Verifies that we have an LLVM binpkg to fall back on.

    Returns:
        True if this function actually created a binpkg, false if one already
        existed.
    """
    if SAVED_LLVM_BINPKG_STAMP.exists():
        pkg = Path(SAVED_LLVM_BINPKG_STAMP.read_text(encoding="utf-8"))
        # Double-check this, since this package is considered a cache artifact
        # by portage. Ergo, it can _technically_ be GC'ed at any time.
        if pkg.exists():
            return False

    pkg = pgo_tools.quickpkg_llvm()
    SAVED_LLVM_BINPKG_STAMP.write_text(str(pkg), encoding="utf-8")
    return True


def restore_llvm_binpkg():
    """Installs the binpkg created by ensure_llvm_binpkg_exists."""
    logging.info("Restoring non-PGO'ed LLVM installation")
    pkg = Path(SAVED_LLVM_BINPKG_STAMP.read_text(encoding="utf-8"))
    assert (
        pkg.exists()
    ), f"Non-PGO'ed binpkg at {pkg} does not exist. Can't restore"
    pgo_tools.run(pgo_tools.generate_quickpkg_restoration_command(pkg))


def find_missing_cross_libs() -> FrozenSet[str]:
    """Returns cross-* libraries that need to be installed for workloads."""
    equery_result = pgo_tools.run(
        ["equery", "l", "--format=$cp", "cross-*/*"],
        check=False,
        stdout=subprocess.PIPE,
    )

    # If no matching package is found, equery will exit with code 3.
    if equery_result.returncode == 3:
        return ALL_NEEDED_CROSS_LIBS

    equery_result.check_returncode()
    has_packages = {x.strip() for x in equery_result.stdout.splitlines()}
    return ALL_NEEDED_CROSS_LIBS - has_packages


def ensure_cross_libs_are_installed():
    """Ensures that we have cross-* libs for all `IMPORTANT_TRIPLES`."""
    missing_packages = find_missing_cross_libs()
    if not missing_packages:
        logging.info("All cross-compiler libraries are already installed")
        return

    missing_packages = sorted(missing_packages)
    logging.info("Installing cross-compiler libs: %s", missing_packages)
    pgo_tools.run(
        ["sudo", "emerge", "-j", "-G"] + missing_packages,
    )


def emerge_pgo_generate_llvm():
    """Emerges a sys-devel/llvm with PGO instrumentation enabled."""
    force_use = (
        "llvm_pgo_generate -llvm_pgo_use"
        # Turn ThinLTO off, since doing so results in way faster builds.
        # This is assumed to be OK, since:
        #   - ThinLTO should have no significant impact on where Clang puts
        #     instrprof counters.
        #   - In practice, both "PGO generated with ThinLTO enabled," and "PGO
        #     generated without ThinLTO enabled," were benchmarked, and the
        #     performance difference between the two was in the noise.
        " -thinlto"
        # Turn ccache off, since if there are valid ccache artifacts from prior
        # runs of this script, ccache will lead to us not getting profdata from
        # those.
        " -wrapper_ccache"
    )
    use = (os.environ.get("USE", "") + " " + force_use).strip()

    # Use FEATURES=ccache since it's not much of a CPU time penalty, and if a
    # user runs this script repeatedly, they'll appreciate it. :)
    force_features = "ccache"
    features = (os.environ.get("FEATURES", "") + " " + force_features).strip()
    logging.info("Building LLVM with USE=%s", shlex.quote(use))
    pgo_tools.run(
        [
            "sudo",
            f"FEATURES={features}",
            f"USE={use}",
            "emerge",
            "sys-devel/llvm",
        ]
    )


def build_profiling_env(profile_dir: Path) -> Dict[str, str]:
    profile_pattern = str(profile_dir / "profile-%m.profraw")
    return {
        "LLVM_PROFILE_OUTPUT_FORMAT": "profraw",
        "LLVM_PROFILE_FILE": profile_pattern,
    }


def ensure_clang_invocations_generate_profiles(clang_bin: str, tmpdir: Path):
    """Raises an exception if clang doesn't generate profraw files.

    Args:
        clang_bin: the path to a clang binary.
        tmpdir: a place where this function can put temporary files.
    """
    tmpdir = tmpdir / "ensure_profiles_generated"
    tmpdir.mkdir(parents=True)
    pgo_tools.run(
        [clang_bin, "--help"],
        extra_env=build_profiling_env(tmpdir),
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    is_empty = next(tmpdir.iterdir(), None) is None
    if is_empty:
        raise ValueError(
            f"The clang binary at {clang_bin} generated no profile"
        )
    shutil.rmtree(tmpdir)


def write_unified_cmake_file(
    into_dir: Path, absl_subdir: Path, gtest_subdir: Path
):
    (into_dir / "CMakeLists.txt").write_text(
        textwrap.dedent(
            f"""\
            cmake_minimum_required(VERSION 3.10)

            project(generate_pgo)

            add_subdirectory({gtest_subdir})
            add_subdirectory({absl_subdir})"""
        ),
        encoding="utf-8",
    )


def fetch_workloads_into(target_dir: Path):
    """Fetches PGO generation workloads into `target_dir`."""
    # The workload here is absl and gtest. The reasoning behind that selection
    # was essentially a mix of:
    # - absl is reasonably-written and self-contained
    # - gtest is needed if tests are to be built; in order to have absl do much
    #   of any linking, gtest is necessary.
    #
    # Use the version of absl that's bundled with ChromeOS at the time of
    # writing.
    target_dir.mkdir(parents=True)

    def fetch_and_extract(gs_url: str, into_dir: Path):
        tgz_full = target_dir / os.path.basename(gs_url)
        pgo_tools.run(
            [
                "gsutil",
                "cp",
                gs_url,
                tgz_full,
            ],
        )
        into_dir.mkdir()

        pgo_tools.run(
            ["tar", "xaf", tgz_full],
            cwd=into_dir,
        )

    absl_dir = target_dir / "absl"
    fetch_and_extract(
        gs_url="gs://chromeos-localmirror/distfiles/"
        "abseil-cpp-a86bb8a97e38bc1361289a786410c0eb5824099c.tar.bz2",
        into_dir=absl_dir,
    )

    gtest_dir = target_dir / "gtest"
    fetch_and_extract(
        gs_url="gs://chromeos-mirror/gentoo/distfiles/"
        "gtest-1b18723e874b256c1e39378c6774a90701d70f7a.tar.gz",
        into_dir=gtest_dir,
    )

    unpacked_absl_dir = read_exactly_one_dirent(absl_dir)
    unpacked_gtest_dir = read_exactly_one_dirent(gtest_dir)
    write_unified_cmake_file(
        into_dir=target_dir,
        absl_subdir=unpacked_absl_dir.relative_to(target_dir),
        gtest_subdir=unpacked_gtest_dir.relative_to(target_dir),
    )


@dataclasses.dataclass(frozen=True)
class WorkloadRunner:
    """Runs benchmark workloads."""

    profraw_dir: Path
    target_dir: Path
    out_dir: Path

    def run(
        self,
        triple: str,
        extra_cflags: Optional[str] = None,
        sysroot: Optional[str] = None,
    ):
        logging.info(
            "Running workload for triple %s, extra cflags %r",
            triple,
            extra_cflags,
        )
        if self.out_dir.exists():
            shutil.rmtree(self.out_dir)
        self.out_dir.mkdir(parents=True)

        clang = triple + "-clang"
        profiling_env = build_profiling_env(self.profraw_dir)
        if sysroot:
            profiling_env["SYSROOT"] = sysroot

        cmake_command: pgo_tools.Command = [
            "cmake",
            "-G",
            "Ninja",
            "-DCMAKE_BUILD_TYPE=RelWithDebInfo",
            f"-DCMAKE_C_COMPILER={clang}",
            f"-DCMAKE_CXX_COMPILER={clang}++",
            "-DABSL_BUILD_TESTING=ON",
            "-DABSL_USE_EXTERNAL_GOOGLETEST=ON",
            "-DABSL_USE_GOOGLETEST_HEAD=OFF",
            "-DABSL_FIND_GOOGLETEST=OFF",
        ]

        if extra_cflags:
            cmake_command += (
                f"-DCMAKE_C_FLAGS={extra_cflags}",
                f"-DCMAKE_CXX_FLAGS={extra_cflags}",
            )

        cmake_command.append(self.target_dir)
        pgo_tools.run(
            cmake_command,
            extra_env=profiling_env,
            cwd=self.out_dir,
        )

        pgo_tools.run(
            ["ninja", "-v", "all"],
            extra_env=profiling_env,
            cwd=self.out_dir,
        )


def read_exactly_one_dirent(directory: Path) -> Path:
    """Returns the single Path under the given directory. Raises otherwise."""
    ents = directory.iterdir()
    ent = next(ents, None)
    if ent is not None:
        if next(ents, None) is None:
            return ent
    raise ValueError(f"Expected exactly one entry under {directory}")


def run_workloads(target_dir: Path) -> Path:
    """Runs all of our workloads in target_dir.

    Args:
        target_dir: a directory that already had `fetch_workloads_into` called
            on it.

    Returns:
        A directory in which profraw files from running the workloads are
        saved.
    """
    profraw_dir = target_dir / "profiles"
    profraw_dir.mkdir()

    out_dir = target_dir / "out"
    runner = WorkloadRunner(
        profraw_dir=profraw_dir,
        target_dir=target_dir,
        out_dir=out_dir,
    )

    # Run the workload once per triple.
    for triple in IMPORTANT_TRIPLES:
        runner.run(
            triple, sysroot=None if triple == HOST_TRIPLE else f"/usr/{triple}"
        )

    # Add a run of ThinLTO, so any ThinLTO-specific lld bits get exercised.
    # Also, since CrOS uses -Os often, exercise that.
    runner.run(HOST_TRIPLE, extra_cflags="-flto=thin -Os")
    return profraw_dir


def convert_profraw_to_pgo_profile(profraw_dir: Path) -> Path:
    """Creates a PGO profile from the profraw profiles in profraw_dir."""
    output = profraw_dir / "merged.prof"
    profile_files = list(profraw_dir.glob("profile-*profraw"))
    if not profile_files:
        raise ValueError("No profraw files generated?")

    logging.info(
        "Creating a PGO profile from %d profraw files", len(profile_files)
    )
    generate_command = [
        "llvm-profdata",
        "merge",
        "--instr",
        f"--output={output}",
    ]
    pgo_tools.run(generate_command + profile_files)
    return output


def main(argv: List[str]):
    logging.basicConfig(
        format=">> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: "
        "%(message)s",
        level=logging.DEBUG,
    )

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--output",
        required=True,
        type=Path,
        help="Where to put the PGO profile",
    )
    parser.add_argument(
        "--use-old-binpkg",
        action="store_true",
        help="""
        This script saves your initial LLVM installation as a binpkg, so it may
        restore that installation later in the build. Passing --use-old-binpkg
        allows this script to use a binpkg from a prior invocation of this
        script.
        """,
    )
    opts = parser.parse_args(argv)

    pgo_tools.exit_if_not_in_chroot()

    output = opts.output

    llvm_binpkg_is_fresh = ensure_llvm_binpkg_exists()
    if not llvm_binpkg_is_fresh and not opts.use_old_binpkg:
        sys.exit(
            textwrap.dedent(
                f"""\
                A LLVM binpkg packed by a previous run of this script is
                available. If you intend this run to be another attempt at the
                previous run, please pass --use-old-binpkg (so the old LLVM
                binpkg is used as our 'baseline'). If you don't, please remove
                the file referring to it at {SAVED_LLVM_BINPKG_STAMP}.
                """
            )
        )

    logging.info("Ensuring `cross-` libraries are installed")
    ensure_cross_libs_are_installed()
    tempdir = Path(tempfile.mkdtemp(prefix="generate_llvm_pgo_profile_"))
    try:
        workloads_path = tempdir / "workloads"
        logging.info("Fetching workloads")
        fetch_workloads_into(workloads_path)

        # If our binpkg is not fresh, we may be operating with a weird LLVM
        # (e.g., a PGO'ed one ;) ). Ensure we always start with that binpkg as
        # our baseline.
        if not llvm_binpkg_is_fresh:
            restore_llvm_binpkg()

        logging.info("Building PGO instrumented LLVM")
        emerge_pgo_generate_llvm()

        logging.info("Ensuring instrumented compilers generate profiles")
        for triple in IMPORTANT_TRIPLES:
            ensure_clang_invocations_generate_profiles(
                triple + "-clang", tempdir
            )

        logging.info("Running workloads")
        profraw_dir = run_workloads(workloads_path)

        # This is a subtle but critical step. The LLVM we're currently working
        # with was built by the LLVM represented _by our binpkg_, which may be
        # a radically different version of LLVM than what was installed (e.g.,
        # it could be from our bootstrap SDK, which could be many months old).
        #
        # If our current LLVM's llvm-profdata is used to interpret the profraw
        # files:
        # 1. The profile generated will be for our new version of clang, and
        #    may therefore be too new for the older version that we still have
        #    to support.
        # 2. There may be silent incompatibilities, as the stability guarantees
        #    of profraw files are not immediately apparent.
        logging.info("Restoring LLVM's binpkg")
        restore_llvm_binpkg()
        pgo_profile = convert_profraw_to_pgo_profile(profraw_dir)
        shutil.copyfile(pgo_profile, output)
    except:
        # Leave the tempdir, as it might help people debug.
        logging.info("NOTE: Tempdir will remain at %s", tempdir)
        raise

    logging.info("Removing now-obsolete tempdir")
    shutil.rmtree(tempdir)
    logging.info("PGO profile is available at %s.", output)


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