Add support for an "inclusion filter" to DexMapper tool

--
MOS_MIGRATED_REVID=123887669

6 files changed, 258 insertions, 11 deletions

diff --git a/src/test/java/com/google/devtools/build/android/ziputils/SplitZipFiltersTest.java b/src/test/java/com/google/devtools/build/android/ziputils/SplitZipFiltersTest.java
new file mode 100644
index 0000000000..7fab1a461c
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/android/ziputils/SplitZipFiltersTest.java
@@ -0,0 +1,63 @@
+// Copyright 2015 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.android.ziputils;
+
+import static com.google.common.truth.Truth.assertThat;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.common.base.Predicate;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
+
+/**
+ * Unit tests for {@link SplitZipFilters}.
+ */
+@RunWith(JUnit4.class)
+public class SplitZipFiltersTest {
+  @Rule public final TemporaryFolder tmp = new TemporaryFolder();
+
+  @Test
+  public void testEntriesIn() throws Exception {
+    File filterZip = createZip("pkg1/test1.class", "pkg2/test2.class");
+    Predicate<String> filter = SplitZipFilters.entriesIn(filterZip.getAbsolutePath());
+    assertThat(filter.apply("pkg1/test1.class")).isTrue();
+    assertThat(filter.apply("pkg2/test1.class")).isFalse();
+    assertThat(filter.apply("pkg2/test2.class")).isTrue();
+    assertThat(filter.apply("pkg1/test1$inner.class")).isFalse();
+  }
+
+  private File createZip(String... entries) throws FileNotFoundException, IOException {
+    File zip = tmp.newFile();
+    int count = 0;
+    try (ZipOutputStream zout = new ZipOutputStream(new FileOutputStream(zip))) {
+      for (String entry : entries) {
+        zout.putNextEntry(new ZipEntry(entry));
+        zout.write(("contents" + count++).getBytes(UTF_8)); // unique content
+        zout.closeEntry();
+      }
+    }
+    return zip;
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/android/ziputils/SplitZipTest.java b/src/test/java/com/google/devtools/build/android/ziputils/SplitZipTest.java
index bba889271e..5d6f861159 100644
--- a/src/test/java/com/google/devtools/build/android/ziputils/SplitZipTest.java
+++ b/src/test/java/com/google/devtools/build/android/ziputils/SplitZipTest.java
@@ -18,6 +18,9 @@ import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import com.google.common.base.Predicates;
+import com.google.common.collect.ImmutableSet;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -426,6 +429,110 @@ public class SplitZipTest {
   }
 
   @Test
+  public void testInputFilter() throws Exception {
+    new ZipFileBuilder()
+        .add("pkg/test.txt", "hello world")
+        .add("pkg/test2.txt", "how are you")
+        .add("pkg/test.class", "hello world")
+        .add("pkg/test2.class", "how are you")
+        .add("pkg/R$attr.class", "bye bye")
+        .create("input.zip");
+
+    new ZipFileBuilder()
+        .add("pkg/test.txt", "hello world")
+        .add("pkg/test2.class", "how are you")
+        .create("expected.zip");
+    byte[] expectedBytes = fileSystem.toByteArray("expected.zip");
+
+    new SplitZip()
+        .addOutput(new ZipOut(fileSystem.getOutputChannel("out/shard1.jar", false),
+            "out/shard1.jar"))
+        .setVerbose(true)
+        .addInput(new ZipIn(fileSystem.getInputChannel("input.zip"), "input.zip"))
+        .setInputFilter(
+            Predicates.in(ImmutableSet.of("pkg/test.txt", "pkg/test2.class", "pkg2/test.class")))
+        .run()
+        .close();
+
+    byte[] outputBytes = fileSystem.toByteArray("out/shard1.jar");
+    assertThat(outputBytes).isEqualTo(expectedBytes);
+  }
+
+  @Test
+  public void testInputFilter_splitDexedClasses() throws Exception {
+    new ZipFileBuilder()
+        .add("pkg/test.class.dex", "hello world")
+        .add("pkg/test2.class", "how are you")
+        .add("pkg/R$attr.class", "bye bye")
+        .create("input.zip");
+
+    new ZipFileBuilder()
+        .add("pkg/test.class.dex", "hello world")
+        .add("pkg/R$attr.class", "bye bye")
+        .create("expected.zip");
+    byte[] expectedBytes = fileSystem.toByteArray("expected.zip");
+
+    new SplitZip()
+        .addOutput(new ZipOut(fileSystem.getOutputChannel("out/shard1.jar", false),
+            "out/shard1.jar"))
+        .setVerbose(true)
+        .addInput(new ZipIn(fileSystem.getInputChannel("input.zip"), "input.zip"))
+        .setInputFilter(
+            Predicates.in(ImmutableSet.of("pkg/test.class", "pkg/R$attr.class")))
+        .setSplitDexedClasses(true)
+        .run()
+        .close();
+
+    byte[] outputBytes = fileSystem.toByteArray("out/shard1.jar");
+    assertThat(outputBytes).isEqualTo(expectedBytes);
+  }
+
+  @Test
+  public void testInputFilter_mainDexFilter() throws Exception {
+    new ZipFileBuilder()
+        .add("pkg1/test1.class", "hello world")
+        .add("pkg2/test1.class", "how are you")
+        .add("pkg1/test2.class", "hello world")
+        .add("pkg2/test2.class", "how are you")
+        .add("pkg1/test3.class", "bye bye")
+        .add("pkg2/test3.class", "bye bye")
+        .create("input.jar");
+
+    String classFileList = "pkg1/test1.class\npkg2/test2.class\n";
+    fileSystem.addFile("main_dex_list.txt", classFileList);
+
+    new SplitZip()
+        .addOutput(new ZipOut(fileSystem.getOutputChannel("out/shard1.jar", false),
+            "out/shard1.jar"))
+        .addOutput(new ZipOut(fileSystem.getOutputChannel("out/shard2.jar", false),
+            "out/shard2.jar"))
+        .setVerbose(true)
+        .setMainClassListFile(fileSystem.getInputStream("main_dex_list.txt"))
+        .addInput(new ZipIn(fileSystem.getInputChannel("input.jar"), "input.jar"))
+        .setInputFilter(
+            Predicates.in(
+                ImmutableSet.of("pkg1/test1.class", "pkg2/test1.class", "pkg3/test1.class")))
+        .setSplitDexedClasses(true)
+        .run()
+        .close();
+
+    // 1st shard contains only main dex list classes also in the filter
+    new ZipFileBuilder()
+        .add("pkg1/test1.class", "hello world")
+        .create("expected/shard1.jar");
+
+    new ZipFileBuilder()
+        .add("pkg2/test1.class", "how are you")
+        .create("expected/shard2.jar");
+
+    assertThat(fileSystem.toByteArray("out/shard1.jar"))
+        .isEqualTo(fileSystem.toByteArray("expected/shard1.jar"));
+
+    assertThat(fileSystem.toByteArray("out/shard2.jar"))
+        .isEqualTo(fileSystem.toByteArray("expected/shard2.jar"));
+  }
+
+  @Test
   public void testVerbose() {
     SplitZip instance = new SplitZip();
     instance.setVerbose(true);
diff --git a/src/tools/android/java/com/google/devtools/build/android/ziputils/BUILD b/src/tools/android/java/com/google/devtools/build/android/ziputils/BUILD
index cc578b7397..4f88bf774e 100644
--- a/src/tools/android/java/com/google/devtools/build/android/ziputils/BUILD
+++ b/src/tools/android/java/com/google/devtools/build/android/ziputils/BUILD
@@ -25,6 +25,7 @@ java_library(
             "DexMapper.java",
             "DexReducer.java",
             "SplitZip.java",
+            "SplitZipFilters.java",
             "Splitter.java",
         ],
     ),
@@ -37,12 +38,11 @@ java_library(
 
 java_library(
     name = "splitter_lib",
-    srcs = glob(
-        [
-            "SplitZip.java",
-            "Splitter.java",
-        ],
-    ),
+    srcs = [
+        "SplitZip.java",
+        "SplitZipFilters.java",
+        "Splitter.java",
+    ],
     visibility = ["//visibility:public"],
     deps = [
         ":ziputils_lib",
@@ -61,6 +61,7 @@ java_binary(
     deps = [
         ":splitter_lib",
         "//src/main/java/com/google/devtools/common/options",
+        "//third_party:guava",
     ],
 )
 
@@ -88,5 +89,6 @@ java_binary(
         ":splitter_lib",
         ":ziputils_lib",
         "//src/main/java/com/google/devtools/common/options",
+        "//third_party:guava",
     ],
 )
diff --git a/src/tools/android/java/com/google/devtools/build/android/ziputils/DexMapper.java b/src/tools/android/java/com/google/devtools/build/android/ziputils/DexMapper.java
index 1e788678a6..694a00156a 100644
--- a/src/tools/android/java/com/google/devtools/build/android/ziputils/DexMapper.java
+++ b/src/tools/android/java/com/google/devtools/build/android/ziputils/DexMapper.java
@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.android.ziputils;
 
+import com.google.common.base.Predicate;
+import com.google.common.base.Predicates;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionsBase;
 import com.google.devtools.common.options.OptionsParser;
@@ -42,12 +44,17 @@ public class DexMapper {
     String resourceFile = options.outputResources;
 
     try {
+      Predicate<String> inputFilter = Predicates.alwaysTrue();
+      if (options.inclusionFilterJar != null) {
+        inputFilter = SplitZipFilters.entriesIn(options.inclusionFilterJar);
+      }
       new SplitZip()
           .setVerbose(false)
           .useDefaultEntryDate()
           .setSplitDexedClasses(options.splitDexedClasses)
           .addInputs(inputs)
           .addOutputs(outputs)
+          .setInputFilter(inputFilter)
           .setMainClassListFile(filterFile)
           .setResourceFile(resourceFile)
           .run()
@@ -99,5 +106,11 @@ public class DexMapper {
         defaultValue = "false",
         help = "Split X.class.dex like X.class if true.  Treated as resources if false.")
     public boolean splitDexedClasses;
+
+    @Option(name = "inclusion_filter_jar",
+        defaultValue = "null",
+        help = "Only copy entries that are listed in the given Jar file.  By default, all entries "
+            + "are copied over.")
+    public String inclusionFilterJar;
   }
 }
diff --git a/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java b/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java
index 4fd1f154b3..dabe4dd34b 100644
--- a/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java
+++ b/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java
@@ -25,6 +25,9 @@ import static com.google.devtools.build.android.ziputils.LocalFileHeader.LOCTIM;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Predicate;
+import com.google.common.base.Predicates;
+import com.google.common.collect.Sets;
 
 import java.io.BufferedReader;
 import java.io.File;
@@ -67,6 +70,7 @@ public class SplitZip implements EntryHandler {
   private final Map<String, ZipOut> assignments = new HashMap<>();
   private final Map<String, CentralDirectory> centralDirectories;
   private final Set<String> classes = new TreeSet<>();
+  private Predicate<String> inputFilter = Predicates.alwaysTrue();
 
   /**
    * Creates an un-configured {@code SplitZip} instance.
@@ -257,6 +261,16 @@ public class SplitZip implements EntryHandler {
   }
 
   /**
+   * Set a predicate to only include files with matching filenames in any of the outputs.  <b>Other
+   * zip entries are dropped</b>, regardless of whether they're classes or resources and regardless
+   * of whether they're listed in {@link #setMainClassListFile}.
+   */
+  public SplitZip setInputFilter(Predicate<String> inputFilter) {
+    this.inputFilter = Preconditions.checkNotNull(inputFilter);
+    return this;
+  }
+
+  /**
    * Executes this {@code SplitZip}, reading content from the configured input locations, creating
    * the specified number of archives, in the configured output directory.
    *
@@ -305,7 +319,8 @@ public class SplitZip implements EntryHandler {
     ZipOut out = assignments.remove(normalizedFilename(header.getFilename()));
     if (out == null) {
       // Skip unassigned file; includes a file with the same name as a previously processed one.
-      // This in particular picks the first .class or .dex file encountered for a given class name.
+      // This in particular picks the first .class or .dex file encountered for a given class name
+      // and drops any file not matched by inputFilter.
       return;
     }
     if (dirEntry == null) {
@@ -391,6 +406,9 @@ public class SplitZip implements EntryHandler {
       CentralDirectory cdir = centralDirectories.get(in.getFilename());
       for (DirectoryEntry entry : cdir.list()) {
         String filename = normalizedFilename(entry.getFilename());
+        if (!inputFilter.apply(filename)) {
+          continue;
+        }
         if (filename.endsWith(".class")) {
           // Only pass classes to the splitter, so that it can do the best job
           // possible distributing them across output files.
@@ -402,15 +420,15 @@ public class SplitZip implements EntryHandler {
         }
       }
     }
-    Splitter entryFilter = new Splitter(outputs.size(), classes.size());
+    Splitter splitter = new Splitter(outputs.size(), classes.size());
     if (filter != null) {
       // Assign files in the filter to the first output file.
-      entryFilter.assign(filter);
-      entryFilter.nextShard(); // minimal initial shard
+      splitter.assign(Sets.filter(filter, inputFilter));
+      splitter.nextShard(); // minimal initial shard
     }
     for (String path : classes) {
       // Use normalized filename so the filter file doesn't have to change
-      int assignment = entryFilter.assign(path);
+      int assignment = splitter.assign(path);
       Preconditions.checkState(assignment >= 0 && assignment < zipOuts.length);
       assignments.put(path, zipOuts[assignment]);
     }
diff --git a/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZipFilters.java b/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZipFilters.java
new file mode 100644
index 0000000000..07c38e8889
--- /dev/null
+++ b/src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZipFilters.java
@@ -0,0 +1,44 @@
+// Copyright 2015 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.android.ziputils;
+
+import com.google.common.base.Predicate;
+import com.google.common.base.Predicates;
+import com.google.common.collect.ImmutableSet;
+
+import java.io.FileInputStream;
+import java.io.IOException;
+
+/**
+ * Custom filters intended for {@link SplitZip#setInputFilter}.
+ */
+public class SplitZipFilters {
+
+  /**
+   * Returns a predicate that returns true for filenames contained in the given zip file.
+   */
+  public static Predicate<String> entriesIn(String filterZip) throws IOException {
+    // Aggregate filenames into a set so Predicates.in is efficient
+    ImmutableSet.Builder<String> filenames = ImmutableSet.builder();
+    @SuppressWarnings("resource") // ZipIn takes ownership but isn't Closable
+    ZipIn zip = new ZipIn(new FileInputStream(filterZip).getChannel(), filterZip);
+    for (DirectoryEntry entry : zip.centralDirectory().list()) {
+      filenames.add(entry.getFilename());
+    }
+    return Predicates.in(filenames.build());
+  }
+
+  private SplitZipFilters() {
+  }
+}