Merge all jars in AAR in aar_import, not just classes.jar

See https://github.com/bazelbuild/bazel/issues/1935

--
MOS_MIGRATED_REVID=137202533
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java b/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
index 23ccc4e..bf2685e 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
@@ -28,8 +28,11 @@
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.rules.RuleConfiguredTargetFactory;
 import com.google.devtools.build.lib.rules.java.JavaCommon;
+import com.google.devtools.build.lib.rules.java.JavaHelper;
 import com.google.devtools.build.lib.rules.java.JavaRuleOutputJarsProvider;
 import com.google.devtools.build.lib.rules.java.JavaRuleOutputJarsProvider.OutputJar;
+import com.google.devtools.build.lib.rules.java.JavaToolchainProvider;
+import com.google.devtools.build.lib.rules.java.Jvm;
 import com.google.devtools.build.lib.vfs.PathFragment;
 
 /**
@@ -37,14 +40,13 @@
  *
  * AAR files are zip archives that contain an Android Manifest, JARs, resources, assets, native
  * libraries, Proguard configuration and lint jars. Currently the aar_import rule supports AARs with
- * an AndroidManifest.xml, classes.jar, res/ and jni/. Assets and additional embedded jars are not
- * yet supported.
+ * an AndroidManifest.xml, classes.jar, libs/, res/ and jni/. Assets are not yet supported.
  *
  * @see <a href="http://tools.android.com/tech-docs/new-build-system/aar-format">AAR Format</a>
  */
 public class AarImport implements RuleConfiguredTargetFactory {
   private static final String ANDROID_MANIFEST = "AndroidManifest.xml";
-  private static final String CLASSES_JAR = "classes.jar";
+  private static final String MERGED_JAR = "classes_and_libs_merged.jar";
 
   @Override
   public ConfiguredTarget create(RuleContext ruleContext)
@@ -52,10 +54,13 @@
     RuleConfiguredTargetBuilder ruleBuilder = new RuleConfiguredTargetBuilder(ruleContext);
     Artifact aar = ruleContext.getPrerequisiteArtifact("aar", Mode.TARGET);
 
-    // classes.jar is required in every AAR.
-    Artifact classesJar = createAarArtifact(ruleContext, CLASSES_JAR);
+    Artifact allAarJars = createAarTreeArtifact(ruleContext, "jars");
+    Artifact jarMergingParams = createAarArtifact(ruleContext, "jar_merging_params");
     ruleContext.registerAction(
-        createEmbeddedJarExtractorActions(ruleContext, aar, CLASSES_JAR, classesJar));
+        createAarEmbeddedJarsExtractorActions(ruleContext, aar, allAarJars, jarMergingParams));
+    Artifact mergedJar = createAarArtifact(ruleContext, MERGED_JAR);
+    ruleContext.registerAction(
+        createAarJarsMergingActions(ruleContext, allAarJars, mergedJar, jarMergingParams));
 
     // AndroidManifest.xml is required in every AAR.
     Artifact androidManifestArtifact = createAarArtifact(ruleContext, ANDROID_MANIFEST);
@@ -66,7 +71,7 @@
     ruleContext.registerAction(
         createManifestExtractorActions(ruleContext, aar, "res/.*", resourcesManifest));
 
-    Artifact resources = createResourcesTreeArtifact(ruleContext);
+    Artifact resources = createAarTreeArtifact(ruleContext, "resources");
     ruleContext.registerAction(
         createManifestFileEntriesExtractorActions(ruleContext, aar, resourcesManifest, resources));
 
@@ -88,13 +93,13 @@
         ruleContext.getImplicitOutputArtifact(AndroidRuleClasses.ANDROID_RESOURCES_ZIP));
 
     NestedSetBuilder<Artifact> filesToBuildBuilder =
-        NestedSetBuilder.<Artifact>stableOrder().add(resources).add(classesJar);
+        NestedSetBuilder.<Artifact>stableOrder().add(resources).add(mergedJar);
 
     Artifact nativeLibs = createAarArtifact(ruleContext, "native_libs.zip");
     ruleContext.registerAction(createAarNativeLibsFilterActions(ruleContext, aar, nativeLibs));
 
     JavaRuleOutputJarsProvider.Builder jarProviderBuilder = new JavaRuleOutputJarsProvider.Builder()
-        .addOutputJar(classesJar, null, null);
+        .addOutputJar(mergedJar, null, null);
     for (TransitiveInfoCollection export : ruleContext.getPrerequisites("exports", Mode.TARGET)) {
       for (OutputJar jar : export.getProvider(JavaRuleOutputJarsProvider.class).getOutputJars()) {
         jarProviderBuilder.addOutputJar(jar);
@@ -145,20 +150,32 @@
         .build(ruleContext);
   }
 
-  // Extracts a jar file from the aar if it exists, otherwise outputs an empty jar file.
-  private static Action[] createEmbeddedJarExtractorActions(RuleContext ruleContext, Artifact aar,
-      String filename, Artifact outputArtifact) {
+  private static Action[] createAarEmbeddedJarsExtractorActions(RuleContext ruleContext,
+      Artifact aar, Artifact jarsTreeArtifact, Artifact singleJarParamFile) {
     return new SpawnAction.Builder()
-        .setExecutable(ruleContext.getExecutablePrerequisite("$aar_embedded_jars_extractor", Mode.HOST))
-        .setMnemonic("AarJarExtractor")
-        .setProgressMessage("Extracting " + filename + " from " + aar.getFilename())
-        .addArgument("--input_archive")
+        .setExecutable(
+            ruleContext.getExecutablePrerequisite("$aar_embedded_jars_extractor", Mode.HOST))
+        .setMnemonic("AarEmbeddedJarsExtractor")
+        .setProgressMessage("Extracting classes.jar and libs/*.jar from " + aar.getFilename())
+        .addArgument("--input_aar")
         .addInputArgument(aar)
-        .addArgument("--filename")
-        .addArgument(filename)
         .addArgument("--output_dir")
-        .addOutput(outputArtifact)
-        .addArgument(outputArtifact.getExecPath().getParentDirectory().getPathString())
+        .addOutputArgument(jarsTreeArtifact)
+        .addArgument("--output_singlejar_param_file")
+        .addOutputArgument(singleJarParamFile)
+        .build(ruleContext);
+  }
+
+  private static Action[] createAarJarsMergingActions(RuleContext ruleContext,
+      Artifact jarsTreeArtifact, Artifact mergedJar, Artifact paramFile) {
+    return singleJarSpawnActionBuilder(ruleContext)
+        .setMnemonic("AarJarsMerger")
+        .setProgressMessage("Merging AAR embedded jars")
+        .addInput(jarsTreeArtifact)
+        .addArgument("--output")
+        .addOutputArgument(mergedJar)
+        .addInput(paramFile)
+        .addArgument("@" + paramFile.getExecPathString())
         .build(ruleContext);
   }
 
@@ -196,8 +213,25 @@
         "_aar", name, ruleContext.getBinOrGenfilesDirectory());
   }
 
-  private static Artifact createResourcesTreeArtifact(RuleContext ruleContext) {
-    PathFragment rootRelativePath = ruleContext.getUniqueDirectory("_aar/unzipped");
+  private static Artifact createAarTreeArtifact(RuleContext ruleContext, String name) {
+    PathFragment rootRelativePath = ruleContext.getUniqueDirectory("_aar/unzipped/" + name);
     return ruleContext.getTreeArtifact(rootRelativePath, ruleContext.getBinOrGenfilesDirectory());
   }
+
+  // Adds the appropriate SpawnAction options depending on if SingleJar is a jar or not.
+  private static SpawnAction.Builder singleJarSpawnActionBuilder(RuleContext ruleContext) {
+    SpawnAction.Builder builder = new SpawnAction.Builder();
+    Artifact singleJar = JavaToolchainProvider.fromRuleContext(ruleContext).getSingleJar();
+    if (singleJar.getFilename().endsWith(".jar")) {
+      builder
+          .setJarExecutable(
+              ruleContext.getHostConfiguration().getFragment(Jvm.class).getJavaExecutable(),
+              singleJar,
+              JavaToolchainProvider.fromRuleContext(ruleContext).getJvmOptions())
+          .addTransitiveInputs(JavaHelper.getHostJavabaseInputs(ruleContext));
+    } else {
+      builder.setExecutable(singleJar);
+    }
+    return builder;
+  }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/rules/android/AarImportTest.java b/src/test/java/com/google/devtools/build/lib/bazel/rules/android/AarImportTest.java
index f87fbe3..0273e6e 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/rules/android/AarImportTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/rules/android/AarImportTest.java
@@ -84,7 +84,7 @@
 
     Artifact resourceTreeArtifact = resourceArtifacts.iterator().next();
     assertThat(resourceTreeArtifact.isTreeArtifact()).isTrue();
-    assertThat(resourceTreeArtifact.getExecPathString()).endsWith("_aar/unzipped/foo");
+    assertThat(resourceTreeArtifact.getExecPathString()).endsWith("_aar/unzipped/resources/foo");
   }
 
   @Test
@@ -131,10 +131,7 @@
     assertThat(outputJars).hasSize(1);
 
     Artifact classesJar = outputJars.iterator().next().getClassJar();
-    assertThat(classesJar.getFilename()).isEqualTo("classes.jar");
-
-    SpawnAction unzipAction = (SpawnAction) getGeneratingAction(classesJar);
-    assertThat(unzipAction.getCommandFilename()).endsWith("aar_embedded_jars_extractor");
+    assertThat(classesJar.getFilename()).isEqualTo("classes_and_libs_merged.jar");
   }
 
   @Test
@@ -151,21 +148,21 @@
   }
 
   @Test
-  public void testExportsPropagatesClassesJar() throws Exception {
+  public void testExportsPropagatesMergedAarJars() throws Exception {
     FileConfiguredTarget appTarget = getFileConfiguredTarget("//java:app.apk");
     Set<Artifact> artifacts = actionsTestUtil().artifactClosureOf(appTarget.getArtifact());
 
     ConfiguredTarget bar = getConfiguredTarget("//a:bar");
     Artifact barClassesJar =
-        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "bar/classes.jar");
-    // Verify that bar/classes.jar was in the artifact closure.
+        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "bar/classes_and_libs_merged.jar");
+    // Verify that bar/classes_and_libs_merged.jar was in the artifact closure.
     assertThat(barClassesJar).isNotNull();
     assertThat(barClassesJar.getArtifactOwner().getLabel()).isEqualTo(bar.getLabel());
 
     ConfiguredTarget foo = getConfiguredTarget("//a:foo");
     Artifact fooClassesJar =
-        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "foo/classes.jar");
-    // Verify that foo/classes.jar was in the artifact closure.
+        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "foo/classes_and_libs_merged.jar");
+    // Verify that foo/classes_and_libs_merged.jar was in the artifact closure.
     assertThat(fooClassesJar).isNotNull();
     assertThat(fooClassesJar.getArtifactOwner().getLabel()).isEqualTo(foo.getLabel());
 
@@ -184,13 +181,13 @@
 
     ConfiguredTarget bar = getConfiguredTarget("//a:bar");
     Artifact barResources =
-        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "_aar/unzipped/bar");
+        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "_aar/unzipped/resources/bar");
     assertThat(barResources).isNotNull();
     assertThat(barResources.getArtifactOwner().getLabel()).isEqualTo(bar.getLabel());
 
     ConfiguredTarget foo = getConfiguredTarget("//a:foo");
     Artifact fooResources =
-        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "_aar/unzipped/foo");
+        ActionsTestUtil.getFirstArtifactEndingWith(artifacts, "_aar/unzipped/resources/foo");
     assertThat(fooResources).isNotNull();
     assertThat(fooResources.getArtifactOwner().getLabel()).isEqualTo(foo.getLabel());
   }
diff --git a/tools/android/aar_embedded_jars_extractor.py b/tools/android/aar_embedded_jars_extractor.py
index 099ebac..22326ef 100644
--- a/tools/android/aar_embedded_jars_extractor.py
+++ b/tools/android/aar_embedded_jars_extractor.py
@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""A tool for extracting jar files from an AAR and failing gracefully.
+"""A tool for extracting all jar files from an AAR.
 
-If the jar file is present within the archive, it is extracted into the output
-directory. If not, an empty jar is created in the output directory.
+An AAR may contain JARs at /classes.jar and /libs/*.jar. This tool extracts all
+of the jars and creates a param file for singlejar to merge them into one jar.
 """
 
 import os
+import re
 import sys
 import zipfile
 
@@ -26,27 +27,30 @@
 
 FLAGS = gflags.FLAGS
 
-gflags.DEFINE_string("input_archive", None, "Input archive")
-gflags.MarkFlagAsRequired("input_archive")
-gflags.DEFINE_string("filename", None, "Filename of JAR to extract")
-gflags.MarkFlagAsRequired("filename")
-gflags.DEFINE_string("output_dir", None, "Output directory")
+gflags.DEFINE_string("input_aar", None, "Input AAR")
+gflags.MarkFlagAsRequired("input_aar")
+gflags.DEFINE_string(
+    "output_singlejar_param_file", None, "Output parameter file for singlejar")
+gflags.MarkFlagAsRequired("output_singlejar_param_file")
+gflags.DEFINE_string("output_dir", None, "Output directory to extract jars in")
 gflags.MarkFlagAsRequired("output_dir")
 
 
-def ExtractEmbeddedJar(input_archive, filename, output_dir):
-  with zipfile.ZipFile(input_archive, "r") as archive:
-    if filename in archive.namelist():
-      archive.extract(filename, output_dir)
-    else:
-      with zipfile.ZipFile(os.path.join(output_dir, filename), "w") as jar:
-        # All jar files must contain META-INF/MANIFEST.MF.
-        jar.writestr("META-INF/MANIFEST.MF", ("Manifest-Version: 1.0\n"
-                                              "Created-By: Bazel\n"))
+def ExtractEmbeddedJars(aar, singlejar_param_file, output_dir):
+  os.makedirs(output_dir)
+  jar_pattern = re.compile("^(classes|libs/.+)\\.jar$")
+  singlejar_param_file.write("--exclude_build_data\n")
+  for name in aar.namelist():
+    if jar_pattern.match(name):
+      singlejar_param_file.write("--sources\n")
+      singlejar_param_file.write(output_dir + "/" + name + "\n")
+      aar.extract(name, output_dir)
 
 
 def main():
-  ExtractEmbeddedJar(FLAGS.input_archive, FLAGS.filename, FLAGS.output_dir)
+  with zipfile.ZipFile(FLAGS.input_aar, "r") as aar:
+    with open(FLAGS.output_singlejar_param_file, "w") as singlejar_param_file:
+      ExtractEmbeddedJars(aar, singlejar_param_file, FLAGS.output_dir)
 
 if __name__ == "__main__":
   FLAGS(sys.argv)
diff --git a/tools/android/aar_embedded_jars_extractor_test.py b/tools/android/aar_embedded_jars_extractor_test.py
index 145637f..9f4dca3 100644
--- a/tools/android/aar_embedded_jars_extractor_test.py
+++ b/tools/android/aar_embedded_jars_extractor_test.py
@@ -14,36 +14,64 @@
 
 """Tests for aar_embedded_jars_extractor."""
 
-import filecmp
 import os
+import shutil
+import StringIO
 import unittest
 import zipfile
 
 from tools.android import aar_embedded_jars_extractor
 
 
-class EmbeddedJarExtractorTest(unittest.TestCase):
+class AarEmbeddedJarsExtractor(unittest.TestCase):
   """Unit tests for aar_embedded_jars_extractor.py."""
 
-  def testPassingJarFile(self):
-    bjar = zipfile.ZipFile("b.jar", "w")
-    bjar.close()
-    azip = zipfile.ZipFile("a.zip", "w")
-    azip.write("b.jar")
-    azip.close()
-    if not os.path.exists("output"):
-      os.mkdir("output")
-    aar_embedded_jars_extractor.ExtractEmbeddedJar("a.zip", "b.jar", "output")
-    self.assertTrue(filecmp.cmp("b.jar", "output/b.jar"))
+  def setUp(self):
+    os.chdir(os.environ["TEST_TMPDIR"])
 
-  def testMissingJarFile(self):
-    azip = zipfile.ZipFile("a.zip", "w")
-    azip.close()
-    if not os.path.exists("output"):
-      os.mkdir("output")
-    aar_embedded_jars_extractor.ExtractEmbeddedJar("a.zip", "b.jar", "output")
-    bjar = zipfile.ZipFile("output/b.jar", "r")
-    self.assertEqual(["META-INF/MANIFEST.MF"], bjar.namelist())
+  def tearDown(self):
+    shutil.rmtree("out_dir")
+
+  def testNoJars(self):
+    aar = zipfile.ZipFile(StringIO.StringIO(), "w")
+    param_file = StringIO.StringIO()
+    aar_embedded_jars_extractor.ExtractEmbeddedJars(aar, param_file, "out_dir")
+    self.assertEqual([], os.listdir("out_dir"))
+    param_file.seek(0)
+    self.assertEqual("--exclude_build_data\n", param_file.read())
+
+  def testClassesJarAndLibsJars(self):
+    aar = zipfile.ZipFile(StringIO.StringIO(), "w")
+    aar.writestr("classes.jar", "")
+    aar.writestr("libs/a.jar", "")
+    aar.writestr("libs/b.jar", "")
+    param_file = StringIO.StringIO()
+    aar_embedded_jars_extractor.ExtractEmbeddedJars(aar, param_file, "out_dir")
+    self.assertItemsEqual(["classes.jar", "libs"], os.listdir("out_dir"))
+    self.assertItemsEqual(["a.jar", "b.jar"], os.listdir("out_dir/libs"))
+    param_file.seek(0)
+    self.assertEqual(
+        ["--exclude_build_data\n",
+         "--sources\n",
+         "out_dir/classes.jar\n",
+         "--sources\n",
+         "out_dir/libs/a.jar\n",
+         "--sources\n",
+         "out_dir/libs/b.jar\n"],
+        param_file.readlines())
+
+  def testOnlyClassesJar(self):
+    aar = zipfile.ZipFile(StringIO.StringIO(), "w")
+    aar.writestr("classes.jar", "")
+    param_file = StringIO.StringIO()
+    aar_embedded_jars_extractor.ExtractEmbeddedJars(aar, param_file, "out_dir")
+    self.assertEqual(["classes.jar"], os.listdir("out_dir"))
+    param_file.seek(0)
+    self.assertEqual(
+        ["--exclude_build_data\n",
+         "--sources\n",
+         "out_dir/classes.jar\n"],
+        param_file.readlines())
 
 
 if __name__ == "__main__":