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

"""Handle most aspects of creating and benchmarking PGO profiles for Rust.

This is meant to be done at Rust uprev time. Ultimately profdata files need
to be placed at

gs://chromeos-localmirror/distfiles/rust-pgo-{rust_version}-frontend.profdata.tz
and
gs://chromeos-localmirror/distfiles/rust-pgo-{rust_version}-llvm.profdata.tz

The intended flow is that you first get the new Rust version in a shape so that
it builds, for instance modifying or adding patches as necessary. Note that if
you need to generate manifests for dev-lang/rust and dev-lang/rust-host before
the profdata files are created, which will cause the `ebuild manifest` command
to fail. One way to handle this is to temporarily delete the lines of the
variable SRC_URI in cros-rustc.eclass which refer to profdata files.

After you have a new working Rust version, you can run the following.

```
$ ./pgo_rust.py generate         # generate profdata files
$ ./pgo_rust.py benchmark-pgo    # benchmark with PGO
$ ./pgo_rust.py benchmark-nopgo  # benchmark without PGO
$ ./pgo_rust.py upload-profdata  # upload profdata to localmirror
```

The benchmark steps aren't strictly necessary, but are recommended and will
upload benchmark data to

gs://chromeos-toolchain-artifacts/rust-pgo/benchmarks/{rust_version}/

Currently by default ripgrep 13.0.0 is used as both the crate to build using an
instrumented Rust while generating profdata, and the crate to build to
benchmark Rust. You may wish to experiment with other crates for either role.
In that case upload your crate to

gs://chromeos-toolchain-artifacts/rust-pgo/crates/{name}-{version}.tar.xz

and use `--crate-name` and `--crate-version` to indicate which crate to build
to generate profdata (or which crate's generated profdata to use), and
`--bench-crate-name` to indicate which crate to build in benchmarks.

Notes on various local and GS locations follow.

Note that currently we need to keep separate profdata files for the LLVM and
frontend components of Rust. This is because LLVM profdata is instrumented by
the system LLVM, but Rust's profdata is instrumented by its own LLVM, which
may have separate profdata.

profdata files accessed by ebuilds must be stored in

gs://chromeos-localmirror/distfiles

Specifically, they go to

gs://chromeos-localmirror/distfiles/rust-pgo-{rust-version}-llvm.profdata.xz

gs://chromeos-localmirror/distfiles/
  rust-pgo-{rust-version}-frontend.profdata.xz

But we can store other data elsewhere, like gs://chromeos-toolchain-artifacts.

GS locations:

{GS_BASE}/crates/ - store crates we may use for generating profiles or
benchmarking PGO optimized Rust compilers

{GS_BASE}/benchmarks/{rust_version}/nopgo/
  {bench_crate_name}-{bench_crate_version}-{triple}

{GS_BASE}/benchmarks/{rust_version}/{crate_name}-{crate_version}/
  {bench_crate_name}-{bench_crate_version}-{triple}

Local locations:

{LOCAL_BASE}/crates/

{LOCAL_BASE}/llvm-profraw/

{LOCAL_BASE}/frontend-profraw/

{LOCAL_BASE}/profdata/{crate_name}-{crate_version}/llvm.profdata

{LOCAL_BASE}/profdata/{crate_name}-{crate_version}/frontend.profdata

{LOCAL_BASE}/benchmarks/{rust_version}/nopgo/
  {bench_crate_name}-{bench_crate_version}-{triple}

{LOCAL_BASE}/benchmarks/{rust_version}/{crate_name}-{crate_version}/
  {bench_crate_name}-{bench_crate_version}-{triple}

{LOCAL_BASE}/llvm.profdata     - must go here to be used by Rust ebuild
{LOCAL_BASE}/frontend.profdata - must go here to be used by Rust ebuild
"""

import argparse
import contextlib
import logging
import os
from pathlib import Path
from pathlib import PurePosixPath
import re
import shutil
import subprocess
import sys
from typing import Dict, List, Optional


TARGET_TRIPLES = [
    "x86_64-cros-linux-gnu",
    "x86_64-pc-linux-gnu",
    "armv7a-cros-linux-gnueabihf",
    "aarch64-cros-linux-gnu",
]

LOCAL_BASE = Path("/tmp/rust-pgo")

GS_BASE = PurePosixPath("/chromeos-toolchain-artifacts/rust-pgo")

GS_DISTFILES = PurePosixPath("/chromeos-localmirror/distfiles")

CRATE_NAME = "ripgrep"

CRATE_VERSION = "13.0.0"


@contextlib.contextmanager
def chdir(new_directory: Path):
    initial_directory = Path.cwd()
    os.chdir(new_directory)
    try:
        yield
    finally:
        os.chdir(initial_directory)


def run(
    args: List,
    *,
    indent: int = 4,
    env: Optional[Dict[str, str]] = None,
    capture_stdout: bool = False,
    message: bool = True,
) -> Optional[str]:
    args = [str(arg) for arg in args]

    if env is None:
        new_env = os.environ
    else:
        new_env = os.environ.copy()
        new_env.update(env)

    if message:
        if env is None:
            logging.info("Running %s", args)
        else:
            logging.info("Running %s in environment %s", args, env)

    result = subprocess.run(
        args,
        env=new_env,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        encoding="utf-8",
        check=False,
    )

    stdout = result.stdout
    stderr = result.stderr
    if indent != 0:
        stdout = re.sub("^", " " * indent, stdout, flags=re.MULTILINE)
        stderr = re.sub("^", " " * indent, stderr, flags=re.MULTILINE)

    if capture_stdout:
        ret = result.stdout
    else:
        logging.info("STDOUT:")
        logging.info(stdout)
        logging.info("STDERR:")
        logging.info(stderr)
        ret = None

    result.check_returncode()

    if message:
        if env is None:
            logging.info("Ran %s\n", args)
        else:
            logging.info("Ran %s in environment %s\n", args, env)

    return ret


def get_rust_version() -> str:
    s = run(["rustc", "--version"], capture_stdout=True)
    m = re.search(r"\d+\.\d+\.\d+", s)
    assert m is not None, repr(s)
    return m.group(0)


def download_unpack_crate(*, crate_name: str, crate_version: str):
    filename_no_extension = f"{crate_name}-{crate_version}"
    gs_path = GS_BASE / "crates" / f"{filename_no_extension}.tar.xz"
    local_path = LOCAL_BASE / "crates"
    shutil.rmtree(
        local_path / f"{crate_name}-{crate_version}", ignore_errors=True
    )
    with chdir(local_path):
        run(["gsutil.py", "cp", f"gs:/{gs_path}", "."])
        run(["xz", "-d", f"{filename_no_extension}.tar.xz"])
        run(["tar", "xvf", f"{filename_no_extension}.tar"])


def build_crate(
    *,
    crate_name: str,
    crate_version: str,
    target_triple: str,
    time_file: Optional[str] = None,
):
    local_path = LOCAL_BASE / "crates" / f"{crate_name}-{crate_version}"
    with chdir(local_path):
        Path(".cargo").mkdir(exist_ok=True)
        with open(".cargo/config.toml", "w") as f:
            f.write(
                "\n".join(
                    (
                        "[source.crates-io]",
                        'replace-with = "vendored-sources"',
                        "",
                        "[source.vendored-sources]",
                        'directory = "vendor"',
                        "",
                        f"[target.{target_triple}]",
                        f'linker = "{target_triple}-clang"',
                        "",
                        "[target.'cfg(all())']",
                        "rustflags = [",
                        '    "-Clto=thin",',
                        '    "-Cembed-bitcode=yes",',
                        "]",
                    )
                )
            )

        run(["cargo", "clean"])

        cargo_cmd = ["cargo", "build", "--release", "--target", target_triple]

        if time_file is None:
            run(cargo_cmd)
        else:
            time_cmd = [
                "/usr/bin/time",
                f"--output={time_file}",
                "--format=wall time (s) %e\nuser time (s) %U\nmax RSS %M\n",
            ]
            run(time_cmd + cargo_cmd)


def build_rust(
    *,
    generate_frontend_profile: bool = False,
    generate_llvm_profile: bool = False,
    use_frontend_profile: bool = False,
    use_llvm_profile: bool = False,
):

    if use_frontend_profile or use_llvm_profile:
        assert (
            not generate_frontend_profile and not generate_llvm_profile
        ), "Can't build a compiler to both use profile information and generate it"

    assert (
        not generate_frontend_profile or not generate_llvm_profile
    ), "Can't generate both frontend and LLVM profile information"

    use = "-rust_profile_frontend_use -rust_profile_llvm_use "
    if generate_frontend_profile:
        use += "rust_profile_frontend_generate "
    if generate_llvm_profile:
        use += "rust_profile_llvm_generate "
    if use_frontend_profile:
        use += "rust_profile_frontend_use_local "
    if use_llvm_profile:
        use += "rust_profile_llvm_use_local "

    # -E to preserve our USE environment variable.
    run(
        ["sudo", "-E", "emerge", "dev-lang/rust", "dev-lang/rust-host"],
        env={"USE": use},
    )


def merge_profdata(llvm_or_frontend, *, source_directory: Path, dest: Path):
    assert llvm_or_frontend in ("llvm", "frontend")

    # The two `llvm-profdata` programs come from different LLVM versions, and may
    # support different versions of the profdata format, so make sure to use the
    # right one.
    llvm_profdata = (
        "/usr/bin/llvm-profdata"
        if llvm_or_frontend == "llvm"
        else "/usr/libexec/rust/llvm-profdata"
    )

    dest.parent.mkdir(parents=True, exist_ok=True)

    files = list(source_directory.glob("*.profraw"))
    run([llvm_profdata, "merge", f"--output={dest}"] + files)


def do_upload_profdata(*, source: Path, dest: PurePosixPath):
    new_path = source.parent / source.name / ".xz"
    run(["xz", "--keep", "--compress", "--force", source])
    upload_file(source=new_path, dest=dest, public_read=True)


def upload_file(
    *, source: Path, dest: PurePosixPath, public_read: bool = False
):
    if public_read:
        run(["gsutil.py", "cp", "-a", "public-read", source, f"gs:/{dest}"])
    else:
        run(["gsutil.py", "cp", source, f"gs:/{dest}"])


def maybe_download_crate(*, crate_name: str, crate_version: str):
    directory = LOCAL_BASE / "crates" / f"{crate_name}-{crate_version}"
    if directory.is_dir():
        logging.info("Crate already downloaded")
    else:
        logging.info("Downloading crate")
        download_unpack_crate(
            crate_name=crate_name, crate_version=crate_version
        )


def generate(args):
    maybe_download_crate(
        crate_name=args.crate_name, crate_version=args.crate_version
    )

    llvm_dir = LOCAL_BASE / "llvm-profraw"
    shutil.rmtree(llvm_dir, ignore_errors=True)
    frontend_dir = LOCAL_BASE / "frontend-profraw"
    shutil.rmtree(frontend_dir, ignore_errors=True)

    logging.info("Building Rust instrumented for llvm")
    build_rust(generate_llvm_profile=True)

    llvm_dir.mkdir(parents=True, exist_ok=True)
    for triple in TARGET_TRIPLES:
        logging.info(
            "Building crate with LLVM instrumentation, for triple %s", triple
        )
        build_crate(
            crate_name=args.crate_name,
            crate_version=args.crate_version,
            target_triple=triple,
        )

    logging.info("Merging LLVM profile data")
    merge_profdata(
        "llvm",
        source_directory=LOCAL_BASE / "llvm-profraw",
        dest=(
            LOCAL_BASE
            / "profdata"
            / f"{args.crate_name}-{args.crate_version}"
            / "llvm.profdata"
        ),
    )

    logging.info("Building Rust instrumented for frontend")
    build_rust(generate_frontend_profile=True)

    frontend_dir.mkdir(parents=True, exist_ok=True)
    for triple in TARGET_TRIPLES:
        logging.info(
            "Building crate with frontend instrumentation, for triple %s",
            triple,
        )
        build_crate(
            crate_name=args.crate_name,
            crate_version=args.crate_version,
            target_triple=triple,
        )

    logging.info("Merging frontend profile data")
    merge_profdata(
        "frontend",
        source_directory=LOCAL_BASE / "frontend-profraw",
        dest=(
            LOCAL_BASE
            / "profdata"
            / f"{args.crate_name}-{args.crate_version}"
            / "frontend.profdata"
        ),
    )


def benchmark_nopgo(args):
    logging.info("Building Rust, no PGO")
    build_rust()

    time_directory = LOCAL_BASE / "benchmarks" / "nopgo"
    logging.info("Benchmarking crate build with no PGO")
    time_directory.mkdir(parents=True, exist_ok=True)
    for triple in TARGET_TRIPLES:
        build_crate(
            crate_name=args.bench_crate_name,
            crate_version=args.bench_crate_version,
            target_triple=triple,
            time_file=(
                time_directory
                / f"{args.bench_crate_name}-{args.bench_crate_version}-{triple}"
            ),
        )

    rust_version = get_rust_version()
    dest_directory = GS_BASE / "benchmarks" / rust_version / "nopgo"
    logging.info("Uploading benchmark data")
    for file in time_directory.iterdir():
        upload_file(
            source=time_directory / file.name, dest=dest_directory / file.name
        )


def benchmark_pgo(args):
    maybe_download_crate(
        crate_name=args.bench_crate_name, crate_version=args.bench_crate_version
    )

    files_dir = Path(
        "/mnt/host/source/src/third_party/chromiumos-overlay",
        "dev-lang/rust/files",
    )

    logging.info("Copying profile data to be used in building Rust")
    run(
        [
            "cp",
            (
                LOCAL_BASE
                / "profdata"
                / f"{args.crate_name}-{args.crate_version}"
                / "llvm.profdata"
            ),
            files_dir,
        ]
    )
    run(
        [
            "cp",
            (
                LOCAL_BASE
                / "profdata"
                / f"{args.crate_name}-{args.crate_version}"
                / "frontend.profdata"
            ),
            files_dir,
        ]
    )

    logging.info("Building Rust with PGO")
    build_rust(use_llvm_profile=True, use_frontend_profile=True)

    time_directory = (
        LOCAL_BASE / "benchmarks" / f"{args.crate_name}-{args.crate_version}"
    )
    time_directory.mkdir(parents=True, exist_ok=True)
    logging.info("Benchmarking crate built with PGO")
    for triple in TARGET_TRIPLES:
        build_crate(
            crate_name=args.bench_crate_name,
            crate_version=args.bench_crate_version,
            target_triple=triple,
            time_file=(
                time_directory
                / f"{args.bench_crate_name}-{args.bench_crate_version}-{triple}"
            ),
        )

    rust_version = get_rust_version()
    dest_directory = (
        GS_BASE
        / "benchmarks"
        / rust_version
        / f"{args.crate_name}-{args.crate_version}"
    )
    logging.info("Uploading benchmark data")
    for file in time_directory.iterdir():
        upload_file(
            source=time_directory / file.name, dest=dest_directory / file.name
        )


def upload_profdata(args):
    directory = (
        LOCAL_BASE / "profdata /" f"{args.crate_name}-{args.crate_version}"
    )
    rust_version = get_rust_version()

    logging.info("Uploading LLVM profdata")
    do_upload_profdata(
        source=directory / "llvm.profdata",
        dest=(GS_DISTFILES / f"rust-pgo-{rust_version}-llvm.profdata.xz"),
    )

    logging.info("Uploading frontend profdata")
    do_upload_profdata(
        source=directory / "frontend.profdata",
        dest=(GS_DISTFILES / f"rust-pgo-{rust_version}-frontend.profdata.xz"),
    )


def main():
    logging.basicConfig(
        stream=sys.stdout, level=logging.NOTSET, format="%(message)s"
    )

    parser = argparse.ArgumentParser(
        prog=sys.argv[0],
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command", help="")
    subparsers.required = True

    parser_generate = subparsers.add_parser(
        "generate",
        help="Generate LLVM and frontend profdata files by building "
        "instrumented Rust compilers, and using them to build the "
        "indicated crate (downloading the crate if necessary).",
    )
    parser_generate.set_defaults(func=generate)
    parser_generate.add_argument(
        "--crate-name", default=CRATE_NAME, help="Name of the crate to build"
    )
    parser_generate.add_argument(
        "--crate-version",
        default=CRATE_VERSION,
        help="Version of the crate to build",
    )

    parser_benchmark_nopgo = subparsers.add_parser(
        "benchmark-nopgo",
        help="Build the Rust compiler without PGO, benchmark "
        "the build of the indicated crate, and upload "
        "the benchmark data.",
    )
    parser_benchmark_nopgo.set_defaults(func=benchmark_nopgo)
    parser_benchmark_nopgo.add_argument(
        "--bench-crate-name",
        default=CRATE_NAME,
        help="Name of the crate whose build to benchmark",
    )
    parser_benchmark_nopgo.add_argument(
        "--bench-crate-version",
        default=CRATE_VERSION,
        help="Version of the crate whose benchmark to build",
    )

    parser_benchmark_pgo = subparsers.add_parser(
        "benchmark-pgo",
        help="Build the Rust compiler using PGO with the indicated "
        "profdata files, benchmark the build of the indicated crate, "
        "and upload the benchmark data.",
    )
    parser_benchmark_pgo.set_defaults(func=benchmark_pgo)
    parser_benchmark_pgo.add_argument(
        "--bench-crate-name",
        default=CRATE_NAME,
        help="Name of the crate whose build to benchmark",
    )
    parser_benchmark_pgo.add_argument(
        "--bench-crate-version",
        default=CRATE_VERSION,
        help="Version of the crate whose benchmark to build",
    )
    parser_benchmark_pgo.add_argument(
        "--crate-name",
        default=CRATE_NAME,
        help="Name of the crate whose profile to use",
    )
    parser_benchmark_pgo.add_argument(
        "--crate-version",
        default=CRATE_VERSION,
        help="Version of the crate whose profile to use",
    )

    parser_upload_profdata = subparsers.add_parser(
        "upload-profdata", help="Upload the profdata files"
    )
    parser_upload_profdata.set_defaults(func=upload_profdata)
    parser_upload_profdata.add_argument(
        "--crate-name",
        default=CRATE_NAME,
        help="Name of the crate whose profile to use",
    )
    parser_upload_profdata.add_argument(
        "--crate-version",
        default=CRATE_VERSION,
        help="Version of the crate whose profile to use",
    )

    args = parser.parse_args()

    (LOCAL_BASE / "crates").mkdir(parents=True, exist_ok=True)
    (LOCAL_BASE / "llvm-profraw").mkdir(parents=True, exist_ok=True)
    (LOCAL_BASE / "frontend-profraw").mkdir(parents=True, exist_ok=True)
    (LOCAL_BASE / "benchmarks").mkdir(parents=True, exist_ok=True)

    args.func(args)

    return 0


if __name__ == "__main__":
    sys.exit(main())