llvm_tools: let `get_llvm_hash` grab llvm-next/llvm hashes

This gives folks a convenient tool to get this info, since it (soon)
won't be in ebuilds.

BUG=b:333462347
TEST=Unittests

Change-Id: I4b9a8e4ed75a459e7bcdfefe5d2e0a713746ce15
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5454855
Reviewed-by: Ryan Beltran <ryanbeltran@chromium.org>
Commit-Queue: George Burgess <gbiv@chromium.org>
Tested-by: George Burgess <gbiv@chromium.org>

4 files changed, 137 insertions, 19 deletions

diff --git a/llvm_tools/get_llvm_hash.py b/llvm_tools/get_llvm_hash.py
index 0922f685..05e14b58 100755
--- a/llvm_tools/get_llvm_hash.py
+++ b/llvm_tools/get_llvm_hash.py
@@ -18,6 +18,8 @@ import tempfile
 from typing import Iterator, Optional, Tuple, Union
 
 import git_llvm_rev
+import llvm_next
+import manifest_utils
 import subprocess_helpers
 
 
@@ -25,7 +27,13 @@ _LLVM_GIT_URL = (
     "https://chromium.googlesource.com/external/github.com/llvm/llvm-project"
 )
 
-KNOWN_HASH_SOURCES = {"google3", "google3-unstable", "tot"}
+KNOWN_HASH_SOURCES = (
+    "google3",
+    "google3-unstable",
+    "llvm",
+    "llvm-next",
+    "tot",
+)
 
 
 def GetVersionFrom(src_dir: Union[Path, str], git_hash: str) -> int:
@@ -355,6 +363,32 @@ def GetLLVMHashAndVersionFromSVNOption(
     return git_hash, version
 
 
+def _FindChromeOSTreeRoot(chromeos_tree_path: Path) -> Path:
+    """Returns the root of a ChromeOS tree, given a path in said tree."""
+    if (chromeos_tree_path / ".repo").exists():
+        return chromeos_tree_path
+
+    for parent in chromeos_tree_path.parents:
+        if (parent / ".repo").exists():
+            return parent
+    raise ValueError(f"{chromeos_tree_path} is not in a repo checkout")
+
+
+def GetCrOSCurrentLLVMHash(chromeos_tree: Path) -> str:
+    """Retrieves the current ChromeOS LLVM hash.
+
+    Args:
+        chromeos_tree: A ChromeOS source tree. This is allowed to be
+        arbitrary subdirectory of an actual ChromeOS tree, for convenience.
+
+    Raises:
+        ManifestValueError if the toolchain manifest doesn't match the
+        expected structure.
+    """
+    chromeos_root = _FindChromeOSTreeRoot(chromeos_tree)
+    return manifest_utils.extract_current_llvm_hash(chromeos_root)
+
+
 class LLVMHash:
     """Provides methods to retrieve a LLVM hash."""
 
@@ -397,15 +431,21 @@ class LLVMHash:
         Returns:
             The hash as a string that corresponds to the LLVM version.
         """
-
         hash_value = GetGitHashFrom(
             GetAndUpdateLLVMProjectInLLVMTools(), version
         )
         return hash_value
 
+    def GetCrOSCurrentLLVMHash(self, chromeos_tree: Path) -> str:
+        """Retrieves the current ChromeOS LLVM hash."""
+        return GetCrOSCurrentLLVMHash(chromeos_tree)
+
+    def GetCrOSLLVMNextHash(self) -> str:
+        """Retrieves the current ChromeOS llvm-next hash."""
+        return llvm_next.LLVM_NEXT_HASH
+
     def GetGoogle3LLVMHash(self) -> str:
         """Retrieves the google3 LLVM hash."""
-
         return self.GetLLVMHash(GetGoogle3LLVMVersion(stable=True))
 
     def GetGoogle3UnstableLLVMHash(self) -> str:
@@ -428,6 +468,7 @@ def main() -> None:
     Parses the command line for the optional command line
     arguments.
     """
+    my_dir = Path(__file__).parent.resolve()
 
     # Create parser and add optional command-line arguments.
     parser = argparse.ArgumentParser(description="Finds the LLVM hash.")
@@ -438,17 +479,36 @@ def main() -> None:
         help="which git hash of LLVM to find. Either a svn revision, or one "
         "of %s" % sorted(KNOWN_HASH_SOURCES),
     )
+    parser.add_argument(
+        "--chromeos_tree",
+        type=Path,
+        required=True,
+        help="""
+        Path to a ChromeOS tree. If not passed, one will be inferred. If none
+        can be inferred, this script will fail.
+        """,
+    )
 
     # Parse command-line arguments.
     args_output = parser.parse_args()
 
     cur_llvm_version = args_output.llvm_version
+    chromeos_tree = args_output.chromeos_tree
+    if not chromeos_tree:
+        # Try to infer this unconditionally, so mishandling of this script can
+        # be more easily detected (which allows more flexibility in the
+        # implementation in the future for things outside of what directly
+        # needs this value).
+        chromeos_tree = _FindChromeOSTreeRoot(my_dir)
 
     new_llvm_hash = LLVMHash()
-
     if isinstance(cur_llvm_version, int):
         # Find the git hash of the specific LLVM version.
         print(new_llvm_hash.GetLLVMHash(cur_llvm_version))
+    elif cur_llvm_version == "llvm":
+        print(new_llvm_hash.GetCrOSCurrentLLVMHash(chromeos_tree))
+    elif cur_llvm_version == "llvm-next":
+        print(new_llvm_hash.GetCrOSLLVMNextHash())
     elif cur_llvm_version == "google3":
         print(new_llvm_hash.GetGoogle3LLVMHash())
     elif cur_llvm_version == "google3-unstable":
diff --git a/llvm_tools/get_llvm_hash_unittest.py b/llvm_tools/get_llvm_hash_unittest.py
index 3fb85861..fda002ce 100755
--- a/llvm_tools/get_llvm_hash_unittest.py
+++ b/llvm_tools/get_llvm_hash_unittest.py
@@ -15,6 +15,7 @@ import unittest
 from unittest import mock
 
 import get_llvm_hash
+import llvm_next
 import subprocess_helpers
 
 
@@ -203,6 +204,12 @@ class TestGetLLVMHash(unittest.TestCase):
 
         self.assertEqual(get_llvm_hash.GetLLVMMajorVersion(), "5432")
 
+    def testGetLLVMNextHash(self):
+        self.assertEqual(
+            get_llvm_hash.LLVMHash().GetCrOSLLVMNextHash(),
+            llvm_next.LLVM_NEXT_HASH,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/llvm_tools/manifest_utils.py b/llvm_tools/manifest_utils.py
index 3d6337bf..e53afa6d 100644
--- a/llvm_tools/manifest_utils.py
+++ b/llvm_tools/manifest_utils.py
@@ -11,7 +11,7 @@ on toolchain projects (llvm-project, etc.) which are public.
 from pathlib import Path
 import shutil
 import subprocess
-from typing import List, Union
+from typing import List, Optional, Union
 from xml.etree import ElementTree
 
 import atomic_write_file
@@ -28,13 +28,64 @@ class UpdateManifestError(Exception):
     """Error occurred when updating the manifest."""
 
 
-def make_xmlparser():
+class ManifestParseError(Exception):
+    """Error occurred when parsing the contents of the manifest."""
+
+
+def make_xmlparser() -> ElementTree.XMLParser:
     """Return a new xmlparser with custom TreeBuilder."""
     return ElementTree.XMLParser(
         target=ElementTree.TreeBuilder(insert_comments=True)
     )
 
 
+def _find_llvm_project_in_manifest_tree(
+    xmlroot: ElementTree.Element,
+) -> Optional[ElementTree.Element]:
+    """Returns the llvm-project `project` in `xmlroot`, if it exists."""
+    for child in xmlroot:
+        if (
+            child.tag == "project"
+            and child.attrib.get("path") == LLVM_PROJECT_PATH
+        ):
+            return child
+    return None
+
+
+def extract_current_llvm_hash(src_tree: Path) -> str:
+    """Returns the current LLVM SHA for the CrOS tree rooted at `src_tree`.
+
+    Raises:
+        ManifestParseError if the manifest didn't have the expected contents.
+    """
+    xmlroot = ElementTree.parse(
+        get_chromeos_manifest_path(src_tree), parser=make_xmlparser()
+    ).getroot()
+    return extract_current_llvm_hash_from_xml(xmlroot)
+
+
+def extract_current_llvm_hash_from_xml(xmlroot: ElementTree.Element) -> str:
+    """Returns the current LLVM SHA for the parsed XML file.
+
+    Raises:
+        ManifestParseError if the manifest didn't have the expected contents.
+    """
+    if xmlroot.tag != "manifest":
+        raise ManifestParseError(
+            f"Root tag is {xmlroot.tag}; should be `manifest`."
+        )
+
+    llvm_project = _find_llvm_project_in_manifest_tree(xmlroot)
+    if llvm_project is None:
+        raise ManifestParseError("No llvm-project `project` found in manifest.")
+
+    revision = llvm_project.attrib.get("revision")
+    if not revision:
+        raise ManifestParseError("Toolchain's `project` has no revision.")
+
+    return revision
+
+
 def update_chromeos_manifest(revision: str, src_tree: Path) -> Path:
     """Replaces the manifest project revision with 'revision'.
 
@@ -73,21 +124,11 @@ def get_chromeos_manifest_path(src_tree: Path) -> Path:
 
 def update_chromeos_manifest_tree(revision: str, xmlroot: ElementTree.Element):
     """Update the revision info for LLVM for a manifest XML root."""
-
-    # This exists mostly for testing.
-    def is_llvm_project(child):
-        return (
-            child.tag == "project" and child.attrib["path"] == LLVM_PROJECT_PATH
-        )
-
-    finder = (child for child in xmlroot if is_llvm_project(child))
-    llvm_project_elem = next(finder, None)
+    llvm_project_elem = _find_llvm_project_in_manifest_tree(xmlroot)
     # Element objects can be falsy, so we need to explicitly check None.
-    if llvm_project_elem is not None:
-        # Update the llvm revision git sha
-        llvm_project_elem.attrib["revision"] = revision
-    else:
+    if llvm_project_elem is None:
         raise UpdateManifestError("xmltree did not have llvm-project")
+    llvm_project_elem.attrib["revision"] = revision
 
 
 def format_manifest(repo_manifest: Path):
diff --git a/llvm_tools/manifest_utils_unittest.py b/llvm_tools/manifest_utils_unittest.py
index 8e283e34..28179413 100755
--- a/llvm_tools/manifest_utils_unittest.py
+++ b/llvm_tools/manifest_utils_unittest.py
@@ -77,6 +77,16 @@ class TestManifestUtils(unittest.TestCase):
         string_root2 = ElementTree.tostring(root)
         self.assertEqual(string_root1, string_root2)
 
+    def test_extract_current_llvm_hash(self):
+        root = ElementTree.fromstring(
+            MANIFEST_FIXTURE,
+            parser=manifest_utils.make_xmlparser(),
+        )
+        self.assertEqual(
+            manifest_utils.extract_current_llvm_hash_from_xml(root),
+            "abcd",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()