xref: /aosp_15_r20/external/toolchain-utils/pgo_tools/update_llvm_manifest.py (revision 760c253c1ed00ce9abd48f8546f08516e57485fe)
1#!/usr/bin/env python3
2# Copyright 2024 The ChromiumOS Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Updates the Manifest file for LLVM.
7
8Often used to pull a new PGO profile in.
9"""
10
11import argparse
12import contextlib
13import logging
14from pathlib import Path
15import re
16import subprocess
17import sys
18from typing import Generator, List
19
20import pgo_tools
21
22
23@contextlib.contextmanager
24def temporarily_add_llvm_next_pgo_to_src_uri(
25    llvm_9999_ebuild: Path,
26) -> Generator[None, None, None]:
27    old_contents = llvm_9999_ebuild.read_text(encoding="utf-8")
28
29    profdata_prefix = "gs://chromeos-localmirror/distfiles/llvm-profdata-"
30    profdata_re = re.compile(
31        # Leave room for a suffix on this, in case we're on the Nth version of
32        # llvm-profdata for some reason.
33        re.escape(profdata_prefix + "${LLVM_HASH}")
34        + r"\S*\.xz\s"
35    )
36    found_urls = list(profdata_re.finditer(old_contents))
37    if len(found_urls) != 1:
38        raise ValueError(
39            f"Want 1 instance of {profdata_re} in {llvm_9999_ebuild}; found "
40            f"{len(found_urls)}"
41        )
42
43    # Insert the new profdata URL right after the old one. The combination of
44    # USE variables gating this file doesn't have to make sense; the URL just
45    # has to be visible to Portage.
46
47    # Note that the regex ended with `\s`, so `.end()` will be after a space.
48    insert_url = profdata_prefix + "${LLVM_NEXT_HASH}.xz "
49    insert_point = found_urls[0].end()
50    new_contents = (
51        old_contents[:insert_point] + insert_url + old_contents[insert_point:]
52    )
53
54    llvm_9999_ebuild.write_text(new_contents, encoding="utf-8")
55    try:
56        yield
57    finally:
58        llvm_9999_ebuild.write_text(old_contents, encoding="utf-8")
59
60
61def update_manifest(llvm_9999_ebuild: Path):
62    subprocess.run(
63        ["ebuild", llvm_9999_ebuild, "manifest"],
64        check=True,
65        stdin=subprocess.DEVNULL,
66    )
67
68
69def main(argv: List[str]) -> None:
70    logging.basicConfig(
71        format=">> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: "
72        "%(message)s",
73        level=logging.INFO,
74    )
75
76    pgo_tools.exit_if_not_in_chroot()
77
78    my_dir = Path(__file__).resolve().parent
79    parser = argparse.ArgumentParser(
80        description=__doc__,
81        formatter_class=argparse.RawDescriptionHelpFormatter,
82    )
83    parser.add_argument(
84        "--llvm-next",
85        action="store_true",
86        help="Also update for the llvm-next PGO profile.",
87    )
88    parser.add_argument(
89        "--chromiumos-overlay",
90        default=my_dir.parent.parent / "chromiumos-overlay",
91        type=Path,
92        help="The chromiumos-overlay directory to work in. Default %(default)s",
93    )
94    opts = parser.parse_args(argv)
95
96    llvm_9999 = opts.chromiumos_overlay / "sys-devel/llvm/llvm-9999.ebuild"
97    if opts.llvm_next:
98        with temporarily_add_llvm_next_pgo_to_src_uri(llvm_9999):
99            update_manifest(llvm_9999)
100    else:
101        update_manifest(llvm_9999)
102
103
104if __name__ == "__main__":
105    main(sys.argv[1:])
106