1 files changed, 85 insertions, 16 deletions
diff --git a/llvm_tools/get_llvm_hash.py b/llvm_tools/get_llvm_hash.py
index 329e8292..83b5ae76 100755
--- a/llvm_tools/get_llvm_hash.py
+++ b/llvm_tools/get_llvm_hash.py
@@ -9,16 +9,18 @@
 from __future__ import print_function
 
 import argparse
+import contextlib
+import functools
 import os
+import re
 import shutil
 import subprocess
 import sys
 import tempfile
-from contextlib import contextmanager
 
 import git_llvm_rev
-from subprocess_helpers import CheckCommand
 from subprocess_helpers import check_output
+from subprocess_helpers import CheckCommand
 
 _LLVM_GIT_URL = ('https://chromium.googlesource.com/external/github.com/llvm'
                  '/llvm-project')
@@ -63,7 +65,66 @@ def GetGitHashFrom(src_dir, version):
       git_llvm_rev.Rev(branch=git_llvm_rev.MAIN_BRANCH, number=version))
 
 
-@contextmanager
+def CheckoutBranch(src_dir, branch):
+  """Checks out and pulls from a branch in a git repo.
+
+  Args:
+    src_dir: The LLVM source tree.
+    branch: The git branch to checkout in src_dir.
+
+  Raises:
+    ValueError: Failed to checkout or pull branch version
+  """
+  CheckCommand(['git', '-C', src_dir, 'checkout', branch])
+  CheckCommand(['git', '-C', src_dir, 'pull'])
+
+
+def ParseLLVMMajorVersion(cmakelist):
+  """Reads CMakeList.txt file contents for LLVMMajor Version.
+
+  Args:
+    cmakelist: contents of CMakeList.txt
+
+  Returns:
+    The major version number as a string
+
+  Raises:
+    ValueError: The major version cannot be parsed from cmakelist
+  """
+  match = re.search(r'\n\s+set\(LLVM_VERSION_MAJOR (?P<major>\d+)\)', cmakelist)
+  if not match:
+    raise ValueError('Failed to parse CMakeList for llvm major version')
+  return match.group('major')
+
+
+@functools.lru_cache(maxsize=1)
+def GetLLVMMajorVersion(git_hash=None):
+  """Reads llvm/CMakeList.txt file contents for LLVMMajor Version.
+
+  Args:
+    git_hash: git hash of llvm version as string or None for top of trunk
+
+  Returns:
+    The major version number as a string
+
+  Raises:
+    ValueError: The major version cannot be parsed from cmakelist or
+      there was a failure to checkout git_hash version
+    FileExistsError: The src directory doe not contain CMakeList.txt
+  """
+  src_dir = GetAndUpdateLLVMProjectInLLVMTools()
+  cmakelists_path = os.path.join(src_dir, 'llvm', 'CMakeLists.txt')
+  if git_hash:
+    CheckCommand(['git', '-C', src_dir, 'checkout', git_hash])
+  try:
+    with open(cmakelists_path) as cmakelists_file:
+      return ParseLLVMMajorVersion(cmakelists_file.read())
+  finally:
+    if git_hash:
+      CheckoutBranch(src_dir, git_llvm_rev.MAIN_BRANCH)
+
+
+@contextlib.contextmanager
 def CreateTempLLVMRepo(temp_dir):
   """Adds a LLVM worktree to 'temp_dir'.
 
@@ -77,7 +138,7 @@ def CreateTempLLVMRepo(temp_dir):
     temp_dir: An absolute path to the temporary directory to put the worktree in
     (obtained via 'tempfile.mkdtemp()').
 
-  Returns:
+  Yields:
     The absolute path to 'temp_dir'.
 
   Raises:
@@ -88,7 +149,8 @@ def CreateTempLLVMRepo(temp_dir):
   abs_path_to_llvm_project_dir = GetAndUpdateLLVMProjectInLLVMTools()
   CheckCommand([
       'git', '-C', abs_path_to_llvm_project_dir, 'worktree', 'add', '--detach',
-      temp_dir, git_llvm_rev.MAIN_BRANCH
+      temp_dir,
+      'origin/%s' % git_llvm_rev.MAIN_BRANCH
   ])
 
   try:
@@ -113,6 +175,9 @@ def GetAndUpdateLLVMProjectInLLVMTools():
   LLVM mirror. In either case, this function will return the absolute path to
   'llvm-project-copy' directory.
 
+  Returns:
+    Absolute path to 'llvm-project-copy' directory in 'llvm_tools'
+
   Raises:
     ValueError: LLVM repo (in 'llvm-project-copy' dir.) has changes or failed to
     checkout to main or failed to fetch from chromium mirror of LLVM.
@@ -125,8 +190,9 @@ def GetAndUpdateLLVMProjectInLLVMTools():
 
   if not os.path.isdir(abs_path_to_llvm_project_dir):
     print(
-        'Checking out LLVM from scratch. This could take a while...\n'
-        '(This should only need to be done once, though.)',
+        (f'Checking out LLVM to {abs_path_to_llvm_project_dir}\n'
+         'so that we can map between commit hashes and revision numbers.\n'
+         'This may take a while, but only has to be done once.'),
         file=sys.stderr)
     os.mkdir(abs_path_to_llvm_project_dir)
 
@@ -142,11 +208,7 @@ def GetAndUpdateLLVMProjectInLLVMTools():
       raise ValueError('LLVM repo in %s has changes, please remove.' %
                        abs_path_to_llvm_project_dir)
 
-    CheckCommand([
-        'git', '-C', abs_path_to_llvm_project_dir, 'checkout',
-        git_llvm_rev.MAIN_BRANCH
-    ])
-    CheckCommand(['git', '-C', abs_path_to_llvm_project_dir, 'pull'])
+    CheckoutBranch(abs_path_to_llvm_project_dir, git_llvm_rev.MAIN_BRANCH)
 
   return abs_path_to_llvm_project_dir
 
@@ -154,6 +216,9 @@ def GetAndUpdateLLVMProjectInLLVMTools():
 def GetGoogle3LLVMVersion(stable):
   """Gets the latest google3 LLVM version.
 
+  Args:
+    stable: boolean, use the stable version or the unstable version
+
   Returns:
     The latest LLVM SVN version as an integer.
 
@@ -178,7 +243,7 @@ def GetGoogle3LLVMVersion(stable):
   return GetVersionFrom(GetAndUpdateLLVMProjectInLLVMTools(), git_hash.rstrip())
 
 
-def is_svn_option(svn_option):
+def IsSvnOption(svn_option):
   """Validates whether the argument (string) is a git hash option.
 
   The argument is used to find the git hash of LLVM.
@@ -186,6 +251,10 @@ def is_svn_option(svn_option):
   Args:
     svn_option: The option passed in as a command line argument.
 
+  Returns:
+    lowercase svn_option if it is a known hash source, otherwise the svn_option
+    as an int
+
   Raises:
     ValueError: Invalid svn option provided.
   """
@@ -212,7 +281,7 @@ def GetLLVMHashAndVersionFromSVNOption(svn_option):
 
   Args:
     svn_option: A valid svn option obtained from the command line.
-      Ex: 'google3', 'tot', or <svn_version> such as 365123.
+      Ex. 'google3', 'tot', or <svn_version> such as 365123.
 
   Returns:
     A tuple that is the LLVM git hash and LLVM version.
@@ -240,7 +309,7 @@ class LLVMHash(object):
   """Provides methods to retrieve a LLVM hash."""
 
   @staticmethod
-  @contextmanager
+  @contextlib.contextmanager
   def CreateTempDirectory():
     temp_dir = tempfile.mkdtemp()
 
@@ -310,7 +379,7 @@ def main():
   parser = argparse.ArgumentParser(description='Finds the LLVM hash.')
   parser.add_argument(
       '--llvm_version',
-      type=is_svn_option,
+      type=IsSvnOption,
       required=True,
       help='which git hash of LLVM to find. Either a svn revision, or one '
       'of %s' % sorted(KNOWN_HASH_SOURCES))