Correctly manage the ProtoApk lifecycle by ensuring it is closed when finished.

Move finding used resources to ProtoResourceUsageAnalyzer for correctness and memory improvements.
New report format: now records the kept resource along with the resource that keep it, as well maintaining the root status for each resource.
  Example line: {number/foo[isRoot: false] = 0x07f...} => [{array/foos[isRoot: true] = 0x...}...]

RELNOTES: None
PiperOrigin-RevId: 208529350

4 files changed, 115 insertions, 65 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/Aapt2ResourceShrinkingAction.java b/src/tools/android/java/com/google/devtools/build/android/Aapt2ResourceShrinkingAction.java
index e9fa1cb338..828981afa5 100644
--- a/src/tools/android/java/com/google/devtools/build/android/Aapt2ResourceShrinkingAction.java
+++ b/src/tools/android/java/com/google/devtools/build/android/Aapt2ResourceShrinkingAction.java
@@ -20,6 +20,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.devtools.build.android.ResourceShrinkerAction.Options;
+import com.google.devtools.build.android.ResourcesZip.ShrunkProtoApk;
 import com.google.devtools.build.android.aapt2.Aapt2ConfigOptions;
 import com.google.devtools.build.android.aapt2.CompiledResources;
 import com.google.devtools.build.android.aapt2.ResourceCompiler;
@@ -108,19 +109,22 @@ public class Aapt2ResourceShrinkingAction {
               .collect(toSet());
 
       if (aapt2ShrinkOptions.useProtoApk) {
-        resourcesZip
-            .shrinkUsingProto(
+        try (final ShrunkProtoApk shrunk =
+            resourcesZip.shrinkUsingProto(
                 packages,
                 options.shrunkJar,
                 options.proguardMapping,
                 options.log,
-                scopedTmp.subDirectoryOf("shrunk-resources"))
-            .writeBinaryTo(linker, options.shrunkApk, aapt2ConfigOptions.resourceTableAsProto)
-            .writeReportTo(options.log)
-            .writeResourcesToZip(options.shrunkResources);
-        if (options.rTxtOutput != null) {
-          // Fufill the contract -- however, we do not generate an R.txt from the shrunk resources.
-          Files.copy(options.rTxt, options.rTxtOutput);
+                scopedTmp.subDirectoryOf("shrunk-resources"))) {
+          shrunk
+              .writeBinaryTo(linker, options.shrunkApk, aapt2ConfigOptions.resourceTableAsProto)
+              .writeReportTo(options.log)
+              .writeResourcesToZip(options.shrunkResources);
+          if (options.rTxtOutput != null) {
+            // Fufill the contract -- however, we do not generate an R.txt from the shrunk
+            // resources.
+            Files.copy(options.rTxt, options.rTxtOutput);
+          }
         }
       } else {
         final ResourceCompiler resourceCompiler =
diff --git a/src/tools/android/java/com/google/devtools/build/android/ResourcesZip.java b/src/tools/android/java/com/google/devtools/build/android/ResourcesZip.java
index c4998ce917..71b2daa14c 100644
--- a/src/tools/android/java/com/google/devtools/build/android/ResourcesZip.java
+++ b/src/tools/android/java/com/google/devtools/build/android/ResourcesZip.java
@@ -30,6 +30,7 @@ import com.google.devtools.build.android.aapt2.ProtoResourceUsageAnalyzer;
 import com.google.devtools.build.android.aapt2.ResourceCompiler;
 import com.google.devtools.build.android.aapt2.ResourceLinker;
 import com.google.devtools.build.android.proto.SerializeFormat.ToolAttributes;
+import java.io.Closeable;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -43,6 +44,7 @@ import java.util.concurrent.ExecutionException;
 import java.util.logging.Logger;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
+import javax.annotation.CheckReturnValue;
 import javax.annotation.Nullable;
 import javax.xml.parsers.ParserConfigurationException;
 import org.xml.sax.SAXException;
@@ -185,6 +187,9 @@ public class ResourcesZip {
 
       if (apkWithAssets != null) {
         ZipFile apkZip = new ZipFile(apkWithAssets.toString());
+        if (apkZip.getEntry("assets/") == null) {
+          zip.addEntry("assets/", new byte[0], compress ? ZipEntry.DEFLATED : ZipEntry.STORED);
+        }
         apkZip
             .stream()
             .filter(entry -> entry.getName().startsWith("assets/"))
@@ -196,7 +201,6 @@ public class ResourcesZip {
                     throw new RuntimeException(e);
                   }
                 });
-        zip.addEntry("assets/", new byte[0], compress ? ZipEntry.DEFLATED : ZipEntry.STORED);
       } else if (Files.exists(assetsRoot)) {
         ZipBuilderVisitorWithDirectories visitor =
             new ZipBuilderVisitorWithDirectories(zip, assetsRoot, "assets");
@@ -243,6 +247,21 @@ public class ResourcesZip {
             ImmutableList.of(workingDirectory), ImmutableList.of(assetsRoot), manifest));
   }
 
+  /**
+   * Shrinks the apk using a protocol buffer apk.
+   *
+   * @param packages The packages of the dependencies. Used to analyze the java code for resource
+   *     references.
+   * @param classJar Used to find resource references in java.
+   * @param proguardMapping Mapping used to decode java references.
+   * @param logFile Destination of the resource shrinker log.
+   * @param workingDirectory Temporary directory for intermediate artifacts.
+   * @return A ShrunkProtoApk, which must be closed when finished.
+   * @throws ParserConfigurationException thrown when the xml parsing not possible.
+   * @throws IOException thrown when the filesystem is going pear shaped.
+   * @throws SAXException thrown when the xml parsing goes badly.
+   */
+  @CheckReturnValue
   public ShrunkProtoApk shrinkUsingProto(
       Set<String> packages,
       Path classJar,
@@ -254,14 +273,17 @@ public class ResourcesZip {
     try (final ProtoApk apk = ProtoApk.readFrom(proto)) {
       final Map<String, Set<String>> toolAttributes = toAttributes();
       // record resources and manifest
-      new ProtoResourceUsageAnalyzer(packages, proguardMapping, logFile)
-          .shrink(
+      final ProtoResourceUsageAnalyzer analyzer =
+          new ProtoResourceUsageAnalyzer(packages, proguardMapping, logFile);
+
+      final ProtoApk shrink =
+          analyzer.shrink(
               apk,
               classJar,
               shrunkApkProto,
               toolAttributes.getOrDefault(SdkConstants.ATTR_KEEP, ImmutableSet.of()),
               toolAttributes.getOrDefault(SdkConstants.ATTR_DISCARD, ImmutableSet.of()));
-      return new ShrunkProtoApk(shrunkApkProto, logFile, ids);
+      return new ShrunkProtoApk(shrink, logFile, ids);
     }
   }
 
@@ -274,12 +296,12 @@ public class ResourcesZip {
         .collect(toMap(Entry::getKey, e -> ImmutableSet.copyOf(e.getValue().getValuesList())));
   }
 
-  static class ShrunkProtoApk {
-    private final Path apk;
+  static class ShrunkProtoApk implements Closeable {
+    private final ProtoApk apk;
     private final Path report;
     private final Path ids;
 
-    ShrunkProtoApk(Path apk, Path report, Path ids) {
+    ShrunkProtoApk(ProtoApk apk, Path report, Path ids) {
       this.apk = apk;
       this.report = report;
       this.ids = ids;
@@ -288,7 +310,7 @@ public class ResourcesZip {
     ShrunkProtoApk writeBinaryTo(ResourceLinker linker, Path binaryOut, boolean writeAsProto)
         throws IOException {
       Files.copy(
-          writeAsProto ? apk : linker.link(ProtoApk.readFrom(apk), ids),
+          writeAsProto ? apk.asApkPath() : linker.link(apk, ids),
           binaryOut,
           StandardCopyOption.REPLACE_EXISTING);
       return this;
@@ -301,10 +323,15 @@ public class ResourcesZip {
 
     ShrunkProtoApk writeResourcesToZip(Path resourcesZip) throws IOException {
       try (final ZipBuilder zip = ZipBuilder.createFor(resourcesZip)) {
-        zip.addEntry("apk.pb", Files.readAllBytes(apk), ZipEntry.STORED);
+        zip.addEntry("apk.pb", Files.readAllBytes(apk.asApkPath()), ZipEntry.STORED);
       }
       return this;
     }
+
+    @Override
+    public void close() throws IOException {
+      apk.close();
+    }
   }
 
   static class ShrunkResources {
diff --git a/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoResourceUsageAnalyzer.java b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoResourceUsageAnalyzer.java
index 171bbf3791..a1cd448827 100644
--- a/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoResourceUsageAnalyzer.java
+++ b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoResourceUsageAnalyzer.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.android.aapt2;
 
 import static java.util.stream.Collectors.joining;
+import static java.util.stream.Collectors.toList;
 
 import com.android.build.gradle.tasks.ResourceUsageAnalyzer;
 import com.android.resources.ResourceType;
@@ -24,7 +25,9 @@ import com.android.tools.lint.detector.api.LintUtils;
 import com.android.utils.XmlUtils;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Splitter;
+import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multimap;
 import com.google.devtools.build.android.aapt2.ProtoApk.ManifestVisitor;
 import com.google.devtools.build.android.aapt2.ProtoApk.ReferenceVisitor;
 import com.google.devtools.build.android.aapt2.ProtoApk.ResourcePackageVisitor;
@@ -34,11 +37,16 @@ import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
+import java.util.ArrayDeque;
 import java.util.Collection;
+import java.util.Deque;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 import java.util.logging.Logger;
+import javax.annotation.CheckReturnValue;
+import javax.annotation.Nullable;
 import javax.xml.parsers.ParserConfigurationException;
 import org.w3c.dom.Attr;
 import org.w3c.dom.DOMException;
@@ -48,6 +56,8 @@ import org.xml.sax.SAXException;
 /** A resource usage analyzer tha functions on apks in protocol buffer format. */
 public class ProtoResourceUsageAnalyzer extends ResourceUsageAnalyzer {
 
+  private static final Logger logger = Logger.getLogger(ProtoResourceUsageAnalyzer.class.getName());
+
   public ProtoResourceUsageAnalyzer(Set<String> resourcePackages, Path mapping, Path logFile)
       throws DOMException, ParserConfigurationException {
     super(resourcePackages, null, null, null, mapping, null, logFile);
@@ -72,7 +82,8 @@ public class ProtoResourceUsageAnalyzer extends ResourceUsageAnalyzer {
    * @param keep A list of resource urls to keep, unused or not.
    * @param discard A list of resource urls to always discard.
    */
-  public void shrink(
+  @CheckReturnValue
+  public ProtoApk shrink(
       ProtoApk apk,
       Path classes,
       Path destination,
@@ -107,30 +118,59 @@ public class ProtoResourceUsageAnalyzer extends ResourceUsageAnalyzer {
 
     keepPossiblyReferencedResources();
 
-    dumpReferences();
-
-    // Remove unused.
-    final ImmutableSet<Resource> unused = ImmutableSet.copyOf(model().findUnused());
+    final List<Resource> resources = model().getResources();
 
-    // ResourceUsageAnalyzer uses the logger to generate the report.
-    Logger logger = Logger.getLogger(getClass().getName());
-    unused.forEach(
-        resource ->
-            logger.fine(
-                "Deleted unused file "
-                    + ((resource.locations != null && resource.locations.getFile() != null)
-                        ? resource.locations.getFile().toString()
-                        : "<apk>" + " for resource " + resource)));
+    List<Resource> roots =
+        resources.stream().filter(r -> r.isKeep() || r.isReachable()).collect(toList());
 
-    apk.copy(
+    final Set<Resource> reachable = findReachableResources(roots);
+    return apk.copy(
         destination,
-        (resourceType, name) ->
-            !unused.contains(
-                Preconditions.checkNotNull(
-                    model().getResource(resourceType, name),
-                    "%s/%s was not declared but is copied!",
-                    resourceType,
-                    name)));
+        (resourceType, name) -> reachable.contains(model().getResource(resourceType, name)));
+  }
+
+  private Set<Resource> findReachableResources(List<Resource> roots) {
+    final Multimap<Resource, Resource> referenceLog = HashMultimap.create();
+    Deque<Resource> queue = new ArrayDeque<>(roots);
+    final Set<Resource> reachable = new HashSet<>();
+    while (!queue.isEmpty()) {
+      Resource resource = queue.pop();
+      if (resource.references != null) {
+        resource.references.forEach(
+            r -> {
+              referenceLog.put(r, resource);
+              queue.add(r);
+            });
+      }
+      // if we see it, it is reachable.
+      reachable.add(resource);
+    }
+
+    // dump resource reference map:
+    final StringBuilder keptResourceLog = new StringBuilder();
+    referenceLog
+        .asMap()
+        .forEach(
+            (resource, referencesTo) ->
+                keptResourceLog
+                    .append(printResource(resource))
+                    .append(" => [")
+                    .append(
+                        referencesTo
+                            .stream()
+                            .map(ProtoResourceUsageAnalyzer::printResource)
+                            .collect(joining(", ")))
+                    .append("]\n"));
+
+    logger.fine("Kept resource references:\n" + keptResourceLog);
+
+    return reachable;
+  }
+
+  private static String printResource(Resource res) {
+    return String.format(
+        "{%s[isRoot: %s] = %s}",
+        res.getUrl(), res.isReachable() || res.isKeep(), "0x" + Integer.toHexString(res.value));
   }
 
   private static final class ResourceDeclarationVisitor implements ResourceVisitor {
@@ -138,11 +178,11 @@ public class ProtoResourceUsageAnalyzer extends ResourceUsageAnalyzer {
     private final ResourceShrinkerUsageModel model;
     private final Set<Integer> packageIds = new HashSet<>();
 
-    public ResourceDeclarationVisitor(ResourceShrinkerUsageModel model) {
+    private ResourceDeclarationVisitor(ResourceShrinkerUsageModel model) {
       this.model = model;
     }
 
-    @javax.annotation.Nullable
+    @Nullable
     @Override
     public ManifestVisitor enteringManifest() {
       return null;
diff --git a/src/tools/android/java/com/google/devtools/build/android/aapt2/ResourceLinker.java b/src/tools/android/java/com/google/devtools/build/android/aapt2/ResourceLinker.java
index 32590f2a7e..2bd13aba59 100644
--- a/src/tools/android/java/com/google/devtools/build/android/aapt2/ResourceLinker.java
+++ b/src/tools/android/java/com/google/devtools/build/android/aapt2/ResourceLinker.java
@@ -72,7 +72,7 @@ import java.util.stream.Stream;
 public class ResourceLinker {
 
   private static final Predicate<String> IS_JAR = s -> s.endsWith(".jar");
-  private static final String PROTO_EXTENSION = "pb-apk";
+  private static final String PROTO_EXTENSION = "-pb.apk";
   private static final String BINARY_EXTENSION = "apk";
   private boolean debug;
   private static final Predicate<DirectoryEntry> IS_FLAT_FILE =
@@ -348,27 +348,6 @@ public class ResourceLinker {
     return fileName.substring(0, lastIndex).concat(".").concat(newExtension);
   }
 
-  private Path convertToBinary(ProtoApk apk) {
-    Path apkPath = apk.asApkPath();
-    try {
-      profiler.startTask("convertToBinary");
-      final Path outPath =
-          workingDirectory.resolveSibling(
-              replaceExtension(apkPath.getFileName().toString(), BINARY_EXTENSION));
-      logger.fine(
-          new AaptCommandBuilder(aapt2)
-              .add("convert")
-              .add("-o", outPath)
-              .add("--output-format", "binary")
-              .add(apkPath.toString())
-              .execute("Converting " + apkPath));
-      profiler.recordEndOf("convertToBinary");
-      return outPath;
-    } catch (IOException e) {
-      throw new LinkError(e);
-    }
-  }
-
   private ProtoApk linkProtoApk(
       CompiledResources compiled,
       Path rTxt,
@@ -571,7 +550,7 @@ public class ResourceLinker {
               .add("--stable-ids", resourceIds)
               .add("--manifest", manifest)
               .addRepeated("-I", StaticLibrary.toPathStrings(linkAgainst))
-              .add("-R", convertToBinary(protoApk))
+              .add("-R", protoApk.asApkPath())
               .add("-o", apk.toString())
               .execute(String.format("Re-linking %s", protoApkPath)));
       return combineApks(protoApkPath, apk, working);