Add apk converted to proto and all attributes from CompiledResources to ResourcesZip.

Add new proto format for tool attributes stored in the AndroidDataXml for storing them in the resources.zip.

RELNOTES:None
PiperOrigin-RevId: 206765679

2 files changed, 199 insertions, 17 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/aapt2/PackagedResources.java b/src/tools/android/java/com/google/devtools/build/android/aapt2/PackagedResources.java
index faad2cd868..2d92b4e071 100644
--- a/src/tools/android/java/com/google/devtools/build/android/aapt2/PackagedResources.java
+++ b/src/tools/android/java/com/google/devtools/build/android/aapt2/PackagedResources.java
@@ -14,6 +14,7 @@
 package com.google.devtools.build.android.aapt2;
 
 import com.google.devtools.build.android.AndroidResourceOutputs;
+import com.google.devtools.build.android.ResourcesZip;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -23,47 +24,64 @@ import javax.annotation.Nullable;
 public class PackagedResources {
 
   private final Path apk;
+  private final Path proto;
   private final Path rTxt;
   private final Path proguardConfig;
   private final Path mainDexProguard;
   private final Path javaSourceDirectory;
   private final Path resourceIds;
+  private final Path attributes;
 
   private PackagedResources(
       Path apk,
+      Path proto,
       Path rTxt,
       Path proguardConfig,
       Path mainDexProguard,
       Path javaSourceDirectory,
-      Path resourceIds) {
+      Path resourceIds,
+      Path attributes) {
     this.apk = apk;
+    this.proto = proto;
     this.rTxt = rTxt;
     this.proguardConfig = proguardConfig;
     this.mainDexProguard = mainDexProguard;
     this.javaSourceDirectory = javaSourceDirectory;
     this.resourceIds = resourceIds;
+    this.attributes = attributes;
   }
 
   public static PackagedResources of(
       Path outPath,
+      Path protoPath,
       Path rTxt,
       Path proguardConfig,
       Path mainDexProguard,
       Path javaSourceDirectory,
-      Path resourceIds)
+      Path resourceIds,
+      Path attributes)
       throws IOException {
     return new PackagedResources(
-        outPath, rTxt, proguardConfig, mainDexProguard, javaSourceDirectory, resourceIds);
+        outPath,
+        protoPath,
+        rTxt,
+        proguardConfig,
+        mainDexProguard,
+        javaSourceDirectory,
+        resourceIds,
+        attributes);
   }
 
   public PackagedResources copyPackageTo(Path packagePath) throws IOException {
     return of(
         copy(apk, packagePath),
+        proto,
         rTxt,
         proguardConfig,
         mainDexProguard,
         javaSourceDirectory,
-        resourceIds);
+        resourceIds,
+        attributes);
   }
 
   public PackagedResources copyRTxtTo(Path rOutput) throws IOException {
@@ -72,11 +90,13 @@ public class PackagedResources {
     }
     return new PackagedResources(
         apk,
+        proto,
         copy(rTxt, rOutput),
         proguardConfig,
         mainDexProguard,
         javaSourceDirectory,
-        resourceIds);
+        resourceIds,
+        attributes);
   }
 
   private Path copy(Path from, Path out) throws IOException {
@@ -91,11 +111,13 @@ public class PackagedResources {
     }
     return of(
         apk,
+        proto,
         rTxt,
         copy(proguardConfig, proguardOut),
         mainDexProguard,
         javaSourceDirectory,
-        resourceIds);
+        resourceIds,
+        attributes);
   }
 
   public PackagedResources copyMainDexProguardTo(Path mainDexProguardOut) throws IOException {
@@ -104,11 +126,13 @@ public class PackagedResources {
     }
     return of(
         apk,
+        proto,
         rTxt,
         proguardConfig,
         copy(mainDexProguard, mainDexProguardOut),
         javaSourceDirectory,
-        resourceIds);
+        resourceIds,
+        attributes);
   }
 
   public PackagedResources createSourceJar(@Nullable Path sourceJarPath) throws IOException {
@@ -116,7 +140,12 @@ public class PackagedResources {
       return this;
     }
     AndroidResourceOutputs.createSrcJar(javaSourceDirectory, sourceJarPath, false);
-    return of(apk, rTxt, proguardConfig, mainDexProguard, sourceJarPath, resourceIds);
+    return of(
+        apk, proto, rTxt, proguardConfig, mainDexProguard, sourceJarPath, resourceIds, attributes);
+  }
+
+  public ResourcesZip packageWith(Path resourceRoot) {
+    return ResourcesZip.fromApkWithProto(proto, attributes, resourceRoot, apk, resourceIds);
   }
 
   public Path getResourceIds() {
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 3839a5d94b..9ea3703d42 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
@@ -21,26 +21,39 @@ import com.google.common.base.Joiner;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multimap;
 import com.google.common.collect.Streams;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.devtools.build.android.AaptCommandBuilder;
+import com.google.devtools.build.android.AndroidCompiledDataDeserializer;
+import com.google.devtools.build.android.AndroidDataWritingVisitor;
+import com.google.devtools.build.android.AndroidResourceMerger.MergingException;
 import com.google.devtools.build.android.AndroidResourceOutputs;
+import com.google.devtools.build.android.FullyQualifiedName;
 import com.google.devtools.build.android.Profiler;
 import com.google.devtools.build.android.aapt2.ResourceCompiler.CompiledType;
+import com.google.devtools.build.android.proto.SerializeFormat.ToolAttributes;
+import com.google.devtools.build.android.xml.Namespaces;
 import com.google.devtools.build.android.ziputils.DirectoryEntry;
 import com.google.devtools.build.android.ziputils.ZipIn;
 import com.google.devtools.build.android.ziputils.ZipOut;
+import java.io.BufferedOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.function.Function;
@@ -263,7 +276,7 @@ public class ResourceLinker {
   }
 
   private List<String> compiledResourcesToPaths(
-      CompiledResources compiled, Predicate<DirectoryEntry> shouldKeep) throws IOException {
+      CompiledResources compiled, Predicate<DirectoryEntry> shouldKeep) {
     // Using sequential streams to maintain the overlay order for aapt2.
     return Stream.concat(include.stream(), Stream.of(compiled))
         .sequential()
@@ -323,9 +336,61 @@ public class ResourceLinker {
     R apply(T arg) throws Throwable;
   }
 
+  private String replaceExtension(String fileName, String newExtension) {
+    int lastIndex = fileName.lastIndexOf('.');
+    if (lastIndex == -1) {
+      return fileName.concat(".").concat(newExtension);
+    }
+    return fileName.substring(0, lastIndex).concat(".").concat(newExtension);
+  }
+
+  public Path convertToBinary(Path protoApk) {
+    try {
+      profiler.startTask("convertToBinary");
+      final Path outPath =
+          workingDirectory.resolveSibling(
+              replaceExtension(protoApk.getFileName().toString(), "apk"));
+      logger.fine(
+          new AaptCommandBuilder(aapt2)
+              .add("convert")
+              .add("-o", outPath)
+              .add("--output-format", "binary")
+              .add(protoApk.toString())
+              .execute("Converting " + protoApk));
+      profiler.recordEndOf("convertToBinary");
+      return outPath;
+    } catch (IOException e) {
+      throw new LinkError(e);
+    }
+  }
+
+  public Path optimizeApk(Path apk) {
+    try {
+      profiler.startTask("optimizeApk");
+      final Path outPath =
+          workingDirectory.resolveSibling(
+              replaceExtension(apk.getFileName().toString(), ".optimized.apk"));
+      logger.fine(
+          new AaptCommandBuilder(aapt2)
+              .forBuildToolsVersion(buildToolsVersion)
+              .forVariantType(VariantType.DEFAULT)
+              .add("optimize")
+              .when(Objects.equals(logger.getLevel(), Level.FINE))
+              .thenAdd("-v")
+              .add("-o", outPath)
+              .add(apk.toString())
+              .execute(String.format("Optimizing %s", apk)));
+      return outPath;
+    } catch (IOException e) {
+      throw new LinkError(e);
+    } finally {
+      profiler.recordEndOf("optimizeApk");
+    }
+  }
+
   public PackagedResources link(CompiledResources compiled) {
     try {
-      final Path outPath = workingDirectory.resolve("bin.apk");
+      final Path outPath = workingDirectory.resolve("bin.pb");
       Path rTxt = workingDirectory.resolve("R.txt");
       Path proguardConfig = workingDirectory.resolve("proguard.cfg");
       Path mainDexProguard = workingDirectory.resolve("proguard.maindex.cfg");
@@ -342,8 +407,7 @@ public class ResourceLinker {
               .thenAdd("--no-version-vectors")
               // Turn off namespaced resources
               .add("--no-static-lib-packages")
-              .when(outputAsProto)
-              .thenAdd("--proto-format")
+              .add("--proto-format")
               .when(Objects.equals(logger.getLevel(), Level.FINE))
               .thenAdd("-v")
               .add("--manifest", compiled.getManifest())
@@ -387,13 +451,35 @@ public class ResourceLinker {
               .thenAdd("--proguard-conditional-keep-rules")
               .add("-o", outPath)
               .execute(String.format("Linking %s", compiled.getManifest())));
-      profiler.recordEndOf("fulllink");
-      profiler.startTask("optimize");
+      profiler.recordEndOf("fulllink").startTask("attributes");
+
+      final Path attributes = workingDirectory.resolve("tool.attributes");
+      // extract tool annotations from the compile resources.
+      final ToolProtoWriter writer = new ToolProtoWriter(attributes);
+      Stream.concat(include.stream(), Stream.of(compiled))
+          .parallel()
+          .map(AndroidCompiledDataDeserializer.create()::readAttributes)
+          .map(Map::entrySet)
+          .flatMap(Set::stream)
+          .distinct()
+          .forEach(e -> e.getValue().writeResource((FullyQualifiedName) e.getKey(), writer));
+      writer.flush();
+
+      profiler.recordEndOf("attributes");
       if (densities.size() < 2) {
         return PackagedResources.of(
-            outPath, rTxt, proguardConfig, mainDexProguard, javaSourceDirectory, resourceIds);
+            outputAsProto ? outPath : convertToBinary(outPath), // convert proto to apk
+            outPath,
+            rTxt,
+            proguardConfig,
+            mainDexProguard,
+            javaSourceDirectory,
+            resourceIds,
+            attributes);
       }
-      final Path optimized = workingDirectory.resolve("optimized.apk");
+
+      profiler.startTask("optimize");
+      final Path optimized = workingDirectory.resolve("optimized.pb");
       logger.fine(
           new AaptCommandBuilder(aapt2)
               .forBuildToolsVersion(buildToolsVersion)
@@ -406,8 +492,16 @@ public class ResourceLinker {
               .add(outPath.toString())
               .execute(String.format("Optimizing %s", compiled.getManifest())));
       profiler.recordEndOf("optimize");
+
       return PackagedResources.of(
-          optimized, rTxt, proguardConfig, mainDexProguard, javaSourceDirectory, resourceIds);
+          outputAsProto ? optimized : convertToBinary(optimized),
+          optimized,
+          rTxt,
+          proguardConfig,
+          mainDexProguard,
+          javaSourceDirectory,
+          resourceIds,
+          attributes);
     } catch (IOException e) {
       throw new LinkError(e);
     }
@@ -442,4 +536,63 @@ public class ResourceLinker {
         .add("baseApk", baseApk)
         .toString();
   }
+
+  private static class ToolProtoWriter implements AndroidDataWritingVisitor {
+
+    final Multimap<String, String> attributes = HashMultimap.create();
+    private final Path out;
+
+    ToolProtoWriter(Path out) {
+      this.out = out;
+    }
+
+    @Override
+    public void flush() throws IOException {
+      ToolAttributes.Builder builder = ToolAttributes.newBuilder();
+      for (Entry<String, Collection<String>> entry : attributes.asMap().entrySet()) {
+        builder.putAttributes(
+            entry.getKey(),
+            ToolAttributes.ToolAttributeValues.newBuilder().addAllValues(entry.getValue()).build());
+      }
+      try (OutputStream stream = new BufferedOutputStream(Files.newOutputStream(out))) {
+        builder.build().writeTo(stream);
+      }
+    }
+
+    @Override
+    public Path copyManifest(Path sourceManifest) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void copyAsset(Path source, String relativeDestinationPath) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void copyResource(Path source, String relativeDestinationPath) throws MergingException {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void defineAttribute(FullyQualifiedName fqn, String name, String value) {
+      attributes.put(removeNamespace(name), value);
+    }
+
+    private String removeNamespace(String qualifiedName) {
+      int indexColon = qualifiedName.indexOf(':');
+      if (indexColon == -1) {
+        return qualifiedName;
+      }
+      return qualifiedName.substring(indexColon);
+    }
+
+    @Override
+    public void defineNamespacesFor(FullyQualifiedName fqn, Namespaces namespaces) {}
+
+    @Override
+    public ValueResourceDefinitionMetadata define(FullyQualifiedName fqn) {
+      throw new UnsupportedOperationException();
+    }
+  }
 }