Expose AndroidResourceUsageAnalyzer methods to enable proto-based resource shrinking without tool attributes.

RELNOTES:
PiperOrigin-RevId: 207132534

5 files changed, 322 insertions, 46 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 8d1be26dab..2db283fdca 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
@@ -33,6 +33,7 @@ import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.ShellQuotedParamsFilePreProcessor;
 import java.io.File;
 import java.nio.file.FileSystems;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
@@ -109,15 +110,17 @@ public class Aapt2ResourceShrinkingAction {
         resourcesZip
             .shrinkUsingProto(
                 packages,
-                options.rTxt,
                 options.shrunkJar,
-                options.primaryManifest,
                 options.proguardMapping,
                 options.log,
                 scopedTmp.subDirectoryOf("shrunk-resources"))
-            .writeBinaryTo(linker, options.shrunkApk)
+            .writeBinaryTo(linker, options.shrunkApk, aapt2ConfigOptions.resourceTableAsProto)
             .writeReportTo(options.log)
-            .writeResourceToZip(options.shrunkResources);
+            .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 =
             ResourceCompiler.create(
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 cf482ff030..ee8e57df8f 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
@@ -21,6 +21,8 @@ import com.google.common.io.ByteStreams;
 import com.google.devtools.build.android.AndroidResourceOutputs.ZipBuilder;
 import com.google.devtools.build.android.AndroidResourceOutputs.ZipBuilderVisitorWithDirectories;
 import com.google.devtools.build.android.aapt2.CompiledResources;
+import com.google.devtools.build.android.aapt2.ProtoApk;
+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 java.io.FileOutputStream;
@@ -142,10 +144,13 @@ public class ResourcesZip {
                 throw new RuntimeException(e);
               }
             });
-    return from(
+    return new ResourcesZip(
         Files.createDirectories(workingDirectory.resolve("res")),
         Files.createDirectories(workingDirectory.resolve("assets")),
-        workingDirectory.resolve("ids.txt"));
+        workingDirectory.resolve("ids.txt"),
+        null,
+        workingDirectory.resolve("apk.pb"),
+        workingDirectory.resolve("tools.attributes.pb"));
   }
 
   /**
@@ -233,14 +238,18 @@ public class ResourcesZip {
 
   public ShrunkProtoApk shrinkUsingProto(
       Set<String> packages,
-      Path rTxt,
       Path classJar,
-      Path primaryManifest,
       Path proguardMapping,
       Path logFile,
       Path workingDirectory)
-      throws ParserConfigurationException {
-    throw new UnsupportedOperationException();
+      throws ParserConfigurationException, IOException, SAXException {
+    final Path shrunkApkProto = workingDirectory.resolve("shrunk.apk.pb");
+    try (final ProtoApk apk = ProtoApk.readFrom(proto)) {
+      // record resources and manifest
+      new ProtoResourceUsageAnalyzer(packages, proguardMapping, logFile)
+          .shrink(apk, classJar, shrunkApkProto);
+      return new ShrunkProtoApk(shrunkApkProto, logFile);
+    }
   }
 
   static class ShrunkProtoApk {
@@ -252,8 +261,12 @@ public class ResourcesZip {
       this.report = report;
     }
 
-    ShrunkProtoApk writeBinaryTo(ResourceLinker linker, Path binaryOut) throws IOException {
-      Files.copy(linker.convertToBinary(apk), binaryOut, StandardCopyOption.REPLACE_EXISTING);
+    ShrunkProtoApk writeBinaryTo(ResourceLinker linker, Path binaryOut, boolean writeAsProto)
+        throws IOException {
+      Files.copy(
+          writeAsProto ? apk : linker.optimizeApk(linker.convertToBinary(apk)),
+          binaryOut,
+          StandardCopyOption.REPLACE_EXISTING);
       return this;
     }
 
@@ -262,7 +275,7 @@ public class ResourcesZip {
       return this;
     }
 
-    ShrunkProtoApk writeResourceToZip(Path resourcesZip) throws IOException {
+    ShrunkProtoApk writeResourcesToZip(Path resourcesZip) throws IOException {
       try (final ZipBuilder zip = ZipBuilder.createFor(resourcesZip)) {
         zip.addEntry("apk.pb", Files.readAllBytes(apk), ZipEntry.STORED);
       }
diff --git a/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java
index a5fe4cd854..33bb6e1bba 100644
--- a/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java
+++ b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoApk.java
@@ -54,6 +54,7 @@ import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.BiPredicate;
+import javax.annotation.Nullable;
 
 /**
  * Provides an interface to an apk in proto format. Since the apk is backed by a zip, it is
@@ -169,6 +170,7 @@ public class ProtoApk implements Closeable {
     }
   }
 
+  /** Traverses the resource table and compiled xml resource using the {@link ResourceVisitor}. */
   public <T extends ResourceVisitor> T visitResources(T visitor) throws IOException {
 
     // visit manifest
@@ -184,17 +186,24 @@ public class ProtoApk implements Closeable {
             : ImmutableList.of();
 
     for (Package pkg : resourceTable.getPackageList()) {
-      ResourcePackageVisitor pkgVisitor = visitor.enteringPackage(pkg.getPackageId().getId());
-      for (Resources.Type type : pkg.getTypeList()) {
-        ResourceTypeVisitor typeVisitor =
-            pkgVisitor.enteringResourceType(
-                type.getTypeId().getId(), ResourceType.getEnum(type.getName()));
-        for (Entry entry : type.getEntryList()) {
-          ResourceValueVisitor entryVisitor =
-              typeVisitor.acceptDeclaration(entry.getName(), entry.getEntryId().getId());
-          for (ConfigValue configValue : entry.getConfigValueList()) {
-            if (configValue.hasValue()) {
-              visitValue(entryVisitor, configValue.getValue(), sourcePool);
+      ResourcePackageVisitor pkgVisitor =
+          visitor.enteringPackage(pkg.getPackageId().getId(), pkg.getPackageName());
+      if (pkgVisitor != null) {
+        for (Resources.Type type : pkg.getTypeList()) {
+          ResourceTypeVisitor typeVisitor =
+              pkgVisitor.enteringResourceType(
+                  type.getTypeId().getId(), ResourceType.getEnum(type.getName()));
+          if (typeVisitor != null) {
+            for (Entry entry : type.getEntryList()) {
+              ResourceValueVisitor entryVisitor =
+                  typeVisitor.enteringDeclaration(entry.getName(), entry.getEntryId().getId());
+              if (entryVisitor != null) {
+                for (ConfigValue configValue : entry.getConfigValueList()) {
+                  if (configValue.hasValue()) {
+                    visitValue(entryVisitor, configValue.getValue(), sourcePool);
+                  }
+                }
+              }
             }
           }
         }
@@ -203,7 +212,7 @@ public class ProtoApk implements Closeable {
     return visitor;
   }
 
-  /** Return the underlying uri for this apk. */
+  /** Accessor for the underlying URI of the apk. */
   public URI asApk() {
     return uri.normalize();
   }
@@ -348,6 +357,7 @@ public class ProtoApk implements Closeable {
         .getAttr()
         .getSymbolList()
         .stream()
+        .filter(Symbol::hasName)
         .map(Symbol::getName)
         .forEach(r -> visitReference(entryVisitor, r));
   }
@@ -391,9 +401,13 @@ public class ProtoApk implements Closeable {
     }
   }
 
-  private void visitXmlResource(Path path, ReferenceVisitor sink) {
+  private void visitXmlResource(Path path, ReferenceVisitor visitor) {
+    if (visitor == null) {
+      return;
+    }
+
     try (InputStream in = Files.newInputStream(path)) {
-      visit(XmlNode.parseFrom(in), sink);
+      visit(XmlNode.parseFrom(in), visitor);
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
@@ -402,14 +416,14 @@ public class ProtoApk implements Closeable {
   private void visit(XmlNode node, ReferenceVisitor sink) {
     if (node.hasElement()) {
       final XmlElement element = node.getElement();
-      element
-          .getAttributeList()
-          .stream()
-          .filter(XmlAttribute::hasCompiledItem)
-          .map(XmlAttribute::getCompiledItem)
-          .filter(Item::hasRef)
-          .map(Item::getRef)
-          .forEach(ref -> visitReference(sink, ref));
+      for (XmlAttribute attribute : element.getAttributeList()) {
+        if (attribute.hasCompiledItem() && attribute.getCompiledItem().hasRef()) {
+          visitReference(sink, attribute.getCompiledItem().getRef());
+        }
+        if (attribute.getResourceId() != 0) {
+          sink.accept(attribute.getResourceId());
+        }
+      }
       element.getChildList().forEach(child -> visit(child, sink));
     }
   }
@@ -430,30 +444,35 @@ public class ProtoApk implements Closeable {
   }
 
   /** Provides an entry point to recording declared and referenced resources in the apk. */
-  public interface ResourceVisitor<T extends ResourceVisitor<T>> {
-    /** Called when entering the manifest. */
+  public interface ResourceVisitor {
+    /** Called when entering the manifest. If null, the manifest is not visited. */
+    @Nullable
     ManifestVisitor enteringManifest();
 
-    /** Called when entering a resource package. */
-    ResourcePackageVisitor enteringPackage(int pkgId);
+    /** Called when entering a resource package. If null, the package is not visited. */
+    @Nullable
+    ResourcePackageVisitor enteringPackage(int pkgId, String packageName);
   }
 
   /** Provides a visitor for packages. */
   public interface ResourcePackageVisitor {
-    /** Called when entering the resource types of the package. */
+    /** Called when entering the resource types of the package. If null, the type is not visited. */
+    @Nullable
     ResourceTypeVisitor enteringResourceType(int typeId, ResourceType type);
   }
 
-  /** Visitor for resource types */
+  /** Visitor for resources types */
   public interface ResourceTypeVisitor {
     /**
      * Called for resource declarations.
      *
      * @param name The name of the resource.
      * @param resourceId The id of the resource, without the package and type.
-     * @return A visitor for accepting references to other resources from the declared resource.
+     * @return A visitor for accepting references to other resources from the declared resource. If
+     *     null, the value is not visited.
      */
-    ResourceValueVisitor acceptDeclaration(String name, int resourceId);
+    @Nullable
+    ResourceValueVisitor enteringDeclaration(String name, int resourceId);
   }
 
   /** A manifest specific resource reference visitor. */
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
new file mode 100644
index 0000000000..af81ce0305
--- /dev/null
+++ b/src/tools/android/java/com/google/devtools/build/android/aapt2/ProtoResourceUsageAnalyzer.java
@@ -0,0 +1,242 @@
+// Copyright 2018 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.aapt2;
+
+import com.android.build.gradle.tasks.ResourceUsageAnalyzer;
+import com.android.resources.ResourceType;
+import com.android.tools.lint.checks.ResourceUsageModel;
+import com.android.tools.lint.checks.ResourceUsageModel.Resource;
+import com.android.tools.lint.detector.api.LintUtils;
+import com.google.common.base.Preconditions;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableSet;
+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;
+import com.google.devtools.build.android.aapt2.ProtoApk.ResourceValueVisitor;
+import com.google.devtools.build.android.aapt2.ProtoApk.ResourceVisitor;
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Set;
+import java.util.logging.Logger;
+import javax.xml.parsers.ParserConfigurationException;
+import org.w3c.dom.DOMException;
+
+/** A resource usage analyzer tha functions on apks in protocol buffer format. */
+public class ProtoResourceUsageAnalyzer extends ResourceUsageAnalyzer {
+
+  public ProtoResourceUsageAnalyzer(Set<String> resourcePackages, Path mapping, Path logFile)
+      throws DOMException, ParserConfigurationException {
+    super(resourcePackages, null, null, null, mapping, null, logFile);
+  }
+
+  private static Resource parse(ResourceUsageModel model, String resourceTypeAndName) {
+    final Iterator<String> iterator = Splitter.on('/').split(resourceTypeAndName).iterator();
+    Preconditions.checkArgument(
+        iterator.hasNext(), "%s invalid resource name", resourceTypeAndName);
+    ResourceType resourceType = ResourceType.getEnum(iterator.next());
+    Preconditions.checkArgument(
+        iterator.hasNext(), "%s invalid resource name", resourceTypeAndName);
+    return model.getResource(resourceType, iterator.next());
+  }
+
+  /**
+   * Calculate and removes unused resource from the {@link ProtoApk}.
+   *
+   * @param apk An apk in the aapt2 proto format.
+   * @param classes The associated classes for the apk.
+   * @param destination Where to write the reduced resources.
+   */
+  public void shrink(ProtoApk apk, Path classes, Path destination) throws IOException {
+
+    // record resources and manifest
+    apk.visitResources(
+        // First, collect all declarations using the declaration visitor.
+        // This allows the model to start with a defined set of resources to build the reference
+        // graph on.
+        apk.visitResources(new ResourceDeclarationVisitor(model())).toUsageVisitor());
+
+    recordClassUsages(classes);
+
+    keepPossiblyReferencedResources();
+
+    dumpReferences();
+
+    // Remove unused.
+    final ImmutableSet<Resource> unused = ImmutableSet.copyOf(model().findUnused());
+
+    // 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)));
+
+    apk.copy(
+        destination,
+        (resourceType, name) ->
+            !unused.contains(
+                Preconditions.checkNotNull(
+                    model().getResource(resourceType, name),
+                    "%s/%s was not declared but is copied!",
+                    resourceType,
+                    name)));
+  }
+
+  private static final class ResourceDeclarationVisitor implements ResourceVisitor {
+
+    private final ResourceShrinkerUsageModel model;
+    private final Set<Integer> packageIds = new HashSet<>();
+
+    public ResourceDeclarationVisitor(ResourceShrinkerUsageModel model) {
+      this.model = model;
+    }
+
+    @javax.annotation.Nullable
+    @Override
+    public ManifestVisitor enteringManifest() {
+      return null;
+    }
+
+    @Override
+    public ResourcePackageVisitor enteringPackage(int pkgId, String packageName) {
+      packageIds.add(pkgId);
+      return (typeId, resourceType) ->
+          (name, resourceId) -> {
+            String hexId =
+                String.format(
+                    "0x%s", Integer.toHexString(((pkgId << 24) | (typeId << 16) | resourceId)));
+            model.addDeclaredResource(resourceType, LintUtils.getFieldName(name), hexId, true);
+            // Skip visiting the definition when collecting declarations.
+            return null;
+          };
+    }
+
+    ResourceUsageVisitor toUsageVisitor() {
+      return new ResourceUsageVisitor(model, ImmutableSet.copyOf(packageIds));
+    }
+  }
+
+  private static final class ResourceUsageVisitor implements ResourceVisitor {
+
+    private final ResourceShrinkerUsageModel model;
+    private final ImmutableSet<Integer> packageIds;
+
+    private ResourceUsageVisitor(
+        ResourceShrinkerUsageModel model, ImmutableSet<Integer> packageIds) {
+      this.model = model;
+      this.packageIds = packageIds;
+    }
+
+    @Override
+    public ManifestVisitor enteringManifest() {
+      return new ManifestVisitor() {
+        @Override
+        public void accept(String name) {
+          ResourceUsageModel.markReachable(model.getResourceFromUrl(name));
+        }
+
+        @Override
+        public void accept(int value) {
+          ResourceUsageModel.markReachable(model.getResource(value));
+        }
+      };
+    }
+
+    @Override
+    public ResourcePackageVisitor enteringPackage(int pkgId, String packageName) {
+      return (typeId, resourceType) ->
+          (name, resourceId) ->
+              new ResourceUsageValueVisitor(
+                  model, model.getResource(resourceType, name), packageIds);
+    }
+  }
+
+  private static final class ResourceUsageValueVisitor implements ResourceValueVisitor {
+
+    private final ResourceUsageModel model;
+    private final Resource declaredResource;
+    private final ImmutableSet<Integer> packageIds;
+
+    private ResourceUsageValueVisitor(
+        ResourceUsageModel model, Resource declaredResource, ImmutableSet<Integer> packageIds) {
+      this.model = model;
+      this.declaredResource = declaredResource;
+      this.packageIds = packageIds;
+    }
+
+    @Override
+    public ReferenceVisitor entering(Path path) {
+      declaredResource.addLocation(new File(path.toString()));
+      return this;
+    }
+
+    @Override
+    public void acceptOpaqueFileType(Path path) {
+      try {
+        String pathString = path.toString();
+        if (pathString.endsWith(".js")) {
+          model.tokenizeJs(
+              declaredResource,
+              new String(java.nio.file.Files.readAllBytes(path), StandardCharsets.UTF_8));
+        } else if (pathString.endsWith(".css")) {
+          model.tokenizeCss(
+              declaredResource,
+              new String(java.nio.file.Files.readAllBytes(path), StandardCharsets.UTF_8));
+        } else if (pathString.endsWith(".html")) {
+          model.tokenizeHtml(
+              declaredResource,
+              new String(java.nio.file.Files.readAllBytes(path), StandardCharsets.UTF_8));
+        } else {
+          // Path is a reference to the apk zip -- unpack it before getting a file reference.
+          model.tokenizeUnknownBinary(
+              declaredResource,
+              java.nio.file.Files.copy(
+                      path,
+                      java.nio.file.Files.createTempFile("binary-resource", null),
+                      StandardCopyOption.REPLACE_EXISTING)
+                  .toFile());
+        }
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    @Override
+    public void accept(String name) {
+      parse(model, name).addReference(declaredResource);
+    }
+
+    @Override
+    public void accept(int value) {
+      if (isInDeclaredPackages(value)) { // ignore references outside of scanned packages.
+        declaredResource.addReference(model.getResource(value));
+      }
+    }
+
+    /** Tests if the id is in any of the scanned packages. */
+    private boolean isInDeclaredPackages(int value) {
+      return packageIds.contains(value >> 24);
+    }
+  }
+}
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 2a47c74378..e33196d1a1 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
@@ -214,9 +214,8 @@ public class ResourceLinker {
               .forBuildToolsVersion(buildToolsVersion)
               .forVariantType(VariantType.LIBRARY)
               .add("link")
-              .when(
-                  outputAsProto) // Used for testing: aapt2 does not output static libraries in
-                                 // proto format.
+              .when(outputAsProto) // Used for testing: aapt2 does not output static libraries in
+              // proto format.
               .thenAdd("--proto-format")
               .when(!outputAsProto)
               .thenAdd("--static-lib")