Utility scripts for shrinkers.

Bug: 160420801
Change-Id: I61e475a0221f08ce6ca3d476659b2201e7aea454
diff --git a/tools/keeprule_benchmark.py b/tools/keeprule_benchmark.py
index 7e7f47c..0a7e063 100755
--- a/tools/keeprule_benchmark.py
+++ b/tools/keeprule_benchmark.py
@@ -10,10 +10,11 @@
 import time
 
 import jdk
+import proguard
 import toolhelper
 import utils
 
-SHRINKERS = ['r8', 'pg']
+SHRINKERS = ['r8'] + proguard.getVersions()
 
 INPUT_PROGRAM = utils.PINNED_R8_JAR
 
@@ -153,6 +154,9 @@
     self.size = size
     self.runs = runs
 
+def isPG(shrinker):
+  return proguard.isValidVersion(shrinker)
+
 def shrinker_args(shrinker, keepfile, output):
   if shrinker == 'r8':
     return [
@@ -165,17 +169,16 @@
       '--min-api', '10000',
       '--pg-conf', keepfile,
       ]
-  elif shrinker == 'pg':
-    return [
-      jdk.GetJavaExecutable(),
-      '-jar', utils.PROGUARD_JAR,
+  elif isPG(shrinker):
+    return proguard.getCmd([
       '-injars', INPUT_PROGRAM,
       '-libraryjars', utils.RT_JAR,
       '-outjars', output,
       '-dontwarn', '**',
       '-optimizationpasses', '2',
       '@' + keepfile,
-    ]
+    ],
+    version=shrinker)
   else:
     assert False, "Unexpected shrinker " + shrinker
 
@@ -219,7 +222,7 @@
       subprocess.check_output(cmd)
       t1 = time.time()
       benchmark_runs.append(t1 - t0)
-      if options.shrinker == 'pg':
+      if isPG(options.shrinker):
         dexout = os.path.join(temp, '%s-out%d-dex.jar' % (name, i))
         dex(out, dexout)
         benchmark_size = utils.uncompressed_size(dexout)
diff --git a/tools/proguard.py b/tools/proguard.py
index b1bde54..5984622 100755
--- a/tools/proguard.py
+++ b/tools/proguard.py
@@ -6,22 +6,65 @@
 # Run ProGuard, Google's internal version
 
 from __future__ import print_function
-import jdk
+
 import os
 import subprocess
 import sys
+from exceptions import ValueError
 
+import jdk
 import utils
 
-PROGUARD_JAR = os.path.join(utils.REPO_ROOT, 'third_party', 'proguard',
-    'proguard_internal_159423826', 'ProGuard_deploy.jar')
+# Internal constants, these should not be used outside this script.
+# Use the friendly utility methods below.
+PG_DIR = os.path.join(utils.THIRD_PARTY, 'proguard')
+DEFAULT = 'pg6'
+DEFAULT_ALIAS = 'pg'
+VERSIONS = {
+  'pg5': os.path.join(PG_DIR, 'proguard5.2.1', 'lib', 'proguard.jar'),
+  'pg6': os.path.join(PG_DIR, 'proguard6.0.1', 'lib', 'proguard.jar'),
+  'pg7': os.path.join(PG_DIR, 'proguard-7.0.0', 'lib', 'proguard.jar'),
+  'pg_internal': os.path.join(
+      PG_DIR, 'proguard_internal_159423826', 'ProGuard_deploy.jar'),
+}
+# Add alias for the default version.
+VERSIONS[DEFAULT_ALIAS] = VERSIONS[DEFAULT]
 
-def run(args, track_memory_file = None, stdout=None, stderr=None):
+# Get versions sorted (nice for argument lists)
+def getVersions():
+  versions = list(VERSIONS.keys())
+  versions.sort()
+  return versions
+
+def isValidVersion(version):
+  return version in VERSIONS
+
+def getValidatedVersion(version):
+  if not isValidVersion(version):
+    raise ValueError("Invalid PG version: '%s'" % version)
+  return version
+
+def getJar(version=DEFAULT):
+  return VERSIONS[getValidatedVersion(version)]
+
+def getRetraceJar(version=DEFAULT):
+  if version == 'pg_internal':
+    raise ValueError("No retrace in internal distribution")
+  return getJar().replace('proguard.jar', 'retrace.jar')
+
+def getCmd(args, version=DEFAULT, jvmArgs=None):
+  cmd = []
+  if jvmArgs:
+    cmd.extend(jvmArgs)
+  cmd.extend([jdk.GetJavaExecutable(), '-jar', getJar(version)])
+  cmd.extend(args)
+  return cmd
+
+def run(args, version=DEFAULT, track_memory_file=None, stdout=None, stderr=None):
   cmd = []
   if track_memory_file:
     cmd.extend(['tools/track_memory.sh', track_memory_file])
-  cmd.extend([jdk.GetJavaExecutable(), '-jar', PROGUARD_JAR])
-  cmd.extend(args)
+  cmd.extend(getCmd(args, version))
   utils.PrintCmd(cmd)
   subprocess.call(cmd, stdout=stdout, stderr=stderr)
 
diff --git a/tools/r8lib_size_compare.py b/tools/r8lib_size_compare.py
index 84dd1a2..568f254 100755
--- a/tools/r8lib_size_compare.py
+++ b/tools/r8lib_size_compare.py
@@ -73,13 +73,17 @@
 
   if not is_output_newer(utils.R8_JAR, r8lib_jar):
     r8lib_memory = os.path.join(tmpdir, 'r8lib%s-memory.txt' % inline_suffix)
+    # TODO(b/160420801): The signature of build_r8lib has changed.
     build_r8lib.build_r8lib(
-        output_path=r8lib_jar, output_map=r8lib_map,
-        extra_args=r8lib_args, track_memory_file=r8lib_memory)
+        output_path=r8lib_jar,
+        output_map=r8lib_map,
+        extra_args=r8lib_args,
+        track_memory_file=r8lib_memory)
 
   pg_output = os.path.join(tmpdir, 'r8lib-pg%s.jar' % inline_suffix)
   pg_memory = os.path.join(tmpdir, 'r8lib-pg%s-memory.txt' % inline_suffix)
   pg_map = os.path.join(tmpdir, 'r8lib-pg%s-map.txt' % inline_suffix)
+  # TODO(b/160420801): This must use proguard.* utils once working again.
   pg_args = ['tools/track_memory.sh', pg_memory,
              'third_party/proguard/proguard6.0.2/bin/proguard.sh',
              '@' + pg_config,
diff --git a/tools/retrace_benchmark.py b/tools/retrace_benchmark.py
index 21bb738..a819ea5 100755
--- a/tools/retrace_benchmark.py
+++ b/tools/retrace_benchmark.py
@@ -4,13 +4,14 @@
 # BSD-style license that can be found in the LICENSE file.
 
 import argparse
-import jdk
-import golem
 import os
 import subprocess
 import sys
 import time
-import toolhelper
+
+import golem
+import jdk
+import proguard
 import utils
 
 RETRACERS = ['r8', 'proguard', 'remapper']
@@ -44,13 +45,7 @@
     retracer_args = [
         '-cp', utils.R8LIB_JAR, 'com.android.tools.r8.retrace.Retrace']
   elif options.retracer == 'proguard':
-    retracer_args = ['-jar',
-                     os.path.join(
-                        utils.THIRD_PARTY,
-                        'proguard',
-                        'proguard6.0.1',
-                        'lib',
-                        'retrace.jar')]
+    retracer_args = ['-jar', proguard.getRetraceJar()]
   elif options.retracer == 'remapper':
     retracer_args = ['-jar',
                      os.path.join(
diff --git a/tools/run_proguard_dx_on_app.py b/tools/run_proguard_dx_on_app.py
index 61ac1d7..bb1d72a 100755
--- a/tools/run_proguard_dx_on_app.py
+++ b/tools/run_proguard_dx_on_app.py
@@ -120,6 +120,7 @@
       track_memory_file = join(temp, utils.MEMORY_USE_TMP_FILE)
     proguard.run(
         args,
+        version='pg_internal',
         track_memory_file = track_memory_file,
         stdout=open(os.devnull, 'w'))
     if options.print_memoryuse:
diff --git a/tools/test_helloexample.py b/tools/test_helloexample.py
index abda493..f81e64c 100755
--- a/tools/test_helloexample.py
+++ b/tools/test_helloexample.py
@@ -19,6 +19,7 @@
 
 import golem
 import jdk
+import proguard
 import utils
 
 HELLO_JAR = os.path.join(utils.BUILD, 'test', 'examples', 'hello.jar')
@@ -87,7 +88,7 @@
   parser = argparse.ArgumentParser(
       description = 'Compile a hello world example program')
   parser.add_argument('--tool',
-                      choices = ['d8', 'r8', 'pg'],
+                      choices = ['d8', 'r8'] + proguard.getVersions(),
                       required = True,
                       help = 'Compiler tool to use.')
   parser.add_argument('--output-mode',
@@ -163,19 +164,17 @@
     if output_mode == 'cf':
       cmd.append('--classfile')
     return [cmd]
-  if tool == 'pg':
+  if proguard.isValidVersion(tool):
     # Build PG invokation with additional rules to silence warnings.
     pg_out = output if output_mode == 'cf' \
       else os.path.join(output_dir, 'pgout.zip')
-    cmds = [[
-      jdk.GetJavaExecutable(),
-      '-jar', utils.PROGUARD_JAR,
+    cmds = [proguard.getCmd([
       '-injars', ':'.join([HELLO_JAR] + extra),
       '-libraryjars', lib,
       '-outjars', pg_out,
       '-dontwarn **',
       '@' + rules_file
-    ]]
+    ], version=tool)]
     if output_mode == 'dex':
       cmds.append(
           GetCompilerPrefix('d8', 'dex', output, pg_out, lib, [], noopt))
diff --git a/tools/utils.py b/tools/utils.py
index 16a0696..4a171ad 100644
--- a/tools/utils.py
+++ b/tools/utils.py
@@ -61,12 +61,6 @@
 GENERATED_LICENSE = os.path.join(GENERATED_LICENSE_DIR, 'LICENSE')
 RT_JAR = os.path.join(REPO_ROOT, 'third_party/openjdk/openjdk-rt-1.8/rt.jar')
 R8LIB_KEEP_RULES = os.path.join(REPO_ROOT, 'src/main/keep.txt')
-PROGUARD_JAR = os.path.join(
-    THIRD_PARTY,
-    'proguard',
-    'proguard6.0.1',
-    'lib',
-    'proguard.jar')
 CF_SEGMENTS_TOOL = os.path.join(THIRD_PARTY, 'cf_segments')
 PINNED_R8_JAR = os.path.join(REPO_ROOT, 'third_party/r8/r8.jar')
 PINNED_PGR8_JAR = os.path.join(REPO_ROOT, 'third_party/r8/r8-pg6.0.1.jar')