AAPT2: Allow output artifacts to be filtered.

A new optional flag has been added to allow a list of artifacts that
should be written as output to be provided. If the flag is provided,
only artifacts that have an output name matching an entry in the list
will be processed.

Test: manually ran against an APK with multiple artifacts in the
configuration and confirmed that only the specified artifacts were
written.
Test: Ran all unit tests.

Change-Id: Ia32b19acf1b2ef3711abf13df08dc7b1aa0cf161

3 files changed, 70 insertions, 3 deletions

diff --git a/tools/aapt2/cmd/Optimize.cpp b/tools/aapt2/cmd/Optimize.cpp
index 2bf91a530526..f40a1df47db9 100644
--- a/tools/aapt2/cmd/Optimize.cpp
+++ b/tools/aapt2/cmd/Optimize.cpp
@@ -73,6 +73,10 @@ struct OptimizeOptions {
   TableFlattenerOptions table_flattener_options;
 
   Maybe<PostProcessingConfiguration> configuration;
+
+  // Set of artifacts to keep when generating multi-APK splits. If the list is empty, all artifacts
+  // are kept and will be written as output.
+  std::unordered_set<std::string> kept_artifacts;
 };
 
 class OptimizeContext : public IAaptContext {
@@ -195,9 +199,12 @@ class OptimizeCommand {
 
     if (options_.configuration && options_.output_dir) {
       MultiApkGenerator generator{apk.get(), context_};
-      MultiApkGeneratorOptions generator_options = {options_.output_dir.value(),
-                                                    options_.configuration.value(),
-                                                    options_.table_flattener_options};
+      MultiApkGeneratorOptions generator_options = {
+          options_.output_dir.value(),
+          options_.configuration.value(),
+          options_.table_flattener_options,
+          options_.kept_artifacts,
+      };
       if (!generator.FromBaseApk(generator_options)) {
         return 1;
       }
@@ -306,6 +313,7 @@ int Optimize(const std::vector<StringPiece>& args) {
   Maybe<std::string> target_abis;
   std::vector<std::string> configs;
   std::vector<std::string> split_args;
+  std::unordered_set<std::string> kept_artifacts;
   bool verbose = false;
   bool print_only = false;
   Flags flags =
@@ -335,6 +343,10 @@ int Optimize(const std::vector<StringPiece>& args) {
                             "Split APK.\nSyntax: path/to/output.apk;<config>[,<config>[...]].\n"
                             "On Windows, use a semicolon ';' separator instead.",
                             &split_args)
+          .OptionalFlagList("--keep-artifacts",
+                            "Comma separated list of artifacts to keep. If none are specified,\n"
+                            "all artifacts will be kept.",
+                            &kept_artifacts)
           .OptionalSwitch("--enable-sparse-encoding",
                           "Enables encoding sparse entries using a binary search tree.\n"
                           "This decreases APK size at the cost of resource retrieval performance.",
@@ -414,6 +426,14 @@ int Optimize(const std::vector<StringPiece>& args) {
       return 0;
     }
 
+    if (!kept_artifacts.empty()) {
+      for (const auto& artifact_str : kept_artifacts) {
+        for (const auto& artifact : util::Tokenize(artifact_str, ',')) {
+          options.kept_artifacts.insert(artifact.to_string());
+        }
+      }
+    }
+
     // Since we know that we are going to process the APK (not just print targets), make sure we
     // have somewhere to write them to.
     if (!options.output_dir) {
diff --git a/tools/aapt2/optimize/MultiApkGenerator.cpp b/tools/aapt2/optimize/MultiApkGenerator.cpp
index 3c96344d8602..da3b8792be69 100644
--- a/tools/aapt2/optimize/MultiApkGenerator.cpp
+++ b/tools/aapt2/optimize/MultiApkGenerator.cpp
@@ -137,6 +137,10 @@ bool MultiApkGenerator::FromBaseApk(const MultiApkGeneratorOptions& options) {
   const StringPiece ext = file::GetExtension(apk_name);
   const std::string base_name = apk_name.substr(0, apk_name.rfind(ext.to_string()));
 
+  std::unordered_set<std::string> artifacts_to_keep = options.kept_artifacts;
+  std::unordered_set<std::string> filtered_artifacts;
+  std::unordered_set<std::string> kept_artifacts;
+
   // For now, just write out the stripped APK since ABI splitting doesn't modify anything else.
   for (const Artifact& artifact : config.artifacts) {
     SourcePathDiagnostics diag{{apk_name}, context_->GetDiagnostics()};
@@ -163,6 +167,20 @@ bool MultiApkGenerator::FromBaseApk(const MultiApkGeneratorOptions& options) {
     ContextWrapper wrapped_context{context_};
     wrapped_context.SetSource({artifact_name});
 
+    if (!options.kept_artifacts.empty()) {
+      const auto& it = artifacts_to_keep.find(artifact_name);
+      if (it == artifacts_to_keep.end()) {
+        filtered_artifacts.insert(artifact_name);
+        if (context_->IsVerbose()) {
+          context_->GetDiagnostics()->Note(DiagMessage(artifact_name) << "skipping artifact");
+        }
+        continue;
+      } else {
+        artifacts_to_keep.erase(it);
+        kept_artifacts.insert(artifact_name);
+      }
+    }
+
     std::unique_ptr<ResourceTable> table =
         FilterTable(artifact, config, *apk_->GetResourceTable(), &wrapped_context, &filters);
     if (!table) {
@@ -197,6 +215,30 @@ bool MultiApkGenerator::FromBaseApk(const MultiApkGeneratorOptions& options) {
     }
   }
 
+  // Make sure all of the requested artifacts were valid. If there are any kept artifacts left,
+  // either the config or the command line was wrong.
+  if (!artifacts_to_keep.empty()) {
+    context_->GetDiagnostics()->Error(
+        DiagMessage() << "The configuration and command line to filter artifacts do not match");
+
+    context_->GetDiagnostics()->Error(DiagMessage() << kept_artifacts.size() << " kept:");
+    for (const auto& artifact : kept_artifacts) {
+      context_->GetDiagnostics()->Error(DiagMessage() << "  " << artifact);
+    }
+
+    context_->GetDiagnostics()->Error(DiagMessage() << filtered_artifacts.size() << " filtered:");
+    for (const auto& artifact : filtered_artifacts) {
+      context_->GetDiagnostics()->Error(DiagMessage() << "  " << artifact);
+    }
+
+    context_->GetDiagnostics()->Error(DiagMessage() << artifacts_to_keep.size() << " missing:");
+    for (const auto& artifact : artifacts_to_keep) {
+      context_->GetDiagnostics()->Error(DiagMessage() << "  " << artifact);
+    }
+
+    return false;
+  }
+
   return true;
 }
 
diff --git a/tools/aapt2/optimize/MultiApkGenerator.h b/tools/aapt2/optimize/MultiApkGenerator.h
index c9b4be88685d..dedb610e1caa 100644
--- a/tools/aapt2/optimize/MultiApkGenerator.h
+++ b/tools/aapt2/optimize/MultiApkGenerator.h
@@ -17,6 +17,10 @@
 #ifndef AAPT2_APKSPLITTER_H
 #define AAPT2_APKSPLITTER_H
 
+#include <memory>
+#include <string>
+#include <unordered_set>
+
 #include "Diagnostics.h"
 #include "LoadedApk.h"
 #include "configuration/ConfigurationParser.h"
@@ -27,6 +31,7 @@ struct MultiApkGeneratorOptions {
   std::string out_dir;
   configuration::PostProcessingConfiguration config;
   TableFlattenerOptions table_flattener_options;
+  std::unordered_set<std::string> kept_artifacts;
 };
 
 /**