Allow for multiple preferred densities in the strip command.

Test: Unit tests pass.

Change-Id: I1f27ac8c36ff3489e4c8e4fce7f3d9cb31df6906

5 files changed, 79 insertions, 33 deletions

diff --git a/tools/aapt2/link/Link.cpp b/tools/aapt2/link/Link.cpp
index f7e0f8f55fd4..0501a3bb5fca 100644
--- a/tools/aapt2/link/Link.cpp
+++ b/tools/aapt2/link/Link.cpp
@@ -1634,7 +1634,7 @@ class LinkCommand {
 
     if (options_.static_lib) {
       if (options_.table_splitter_options.config_filter != nullptr ||
-          options_.table_splitter_options.preferred_density) {
+          !options_.table_splitter_options.preferred_densities.empty()) {
         context_->GetDiagnostics()
             ->Warn(DiagMessage()
                    << "can't strip resources when building static library");
@@ -2107,8 +2107,7 @@ int Link(const std::vector<StringPiece>& args) {
                         << "Preferred density must only be a density value");
       return 1;
     }
-    options.table_splitter_options.preferred_density =
-        preferred_density_config.density;
+    options.table_splitter_options.preferred_densities.push_back(preferred_density_config.density);
   }
 
   if (!options.static_lib && stable_id_file_path) {
diff --git a/tools/aapt2/split/TableSplitter.cpp b/tools/aapt2/split/TableSplitter.cpp
index 38cfd2e61054..27e13d81ff49 100644
--- a/tools/aapt2/split/TableSplitter.cpp
+++ b/tools/aapt2/split/TableSplitter.cpp
@@ -19,6 +19,7 @@
 #include <algorithm>
 #include <map>
 #include <set>
+#include <unordered_set>
 #include <unordered_map>
 #include <vector>
 
@@ -124,29 +125,35 @@ class SplitValueSelector {
  * leaving only the preferred density behind.
  */
 static void MarkNonPreferredDensitiesAsClaimed(
-    uint16_t preferred_density, const ConfigDensityGroups& density_groups,
+    const std::vector<uint16_t>& preferred_densities, const ConfigDensityGroups& density_groups,
     ConfigClaimedMap* config_claimed_map) {
   for (auto& entry : density_groups) {
     const ConfigDescription& config = entry.first;
     const std::vector<ResourceConfigValue*>& related_values = entry.second;
 
-    ConfigDescription target_density = config;
-    target_density.density = preferred_density;
-    ResourceConfigValue* best_value = nullptr;
+    // There can be multiple best values if there are multiple preferred densities.
+    std::unordered_set<ResourceConfigValue*> best_values;
+
+    // For each preferred density, find the value that is the best.
+    for (uint16_t preferred_density : preferred_densities) {
+      ConfigDescription target_density = config;
+      target_density.density = preferred_density;
+      ResourceConfigValue* best_value = nullptr;
+      for (ResourceConfigValue* this_value : related_values) {
+        if (!best_value || this_value->config.isBetterThan(best_value->config, &target_density)) {
+          best_value = this_value;
+        }
+      }
+      CHECK(best_value != nullptr);
+      best_values.insert(best_value);
+    }
+
+    // Claim all the values that aren't the best so that they will be removed from the base.
     for (ResourceConfigValue* this_value : related_values) {
-      if (!best_value) {
-        best_value = this_value;
-      } else if (this_value->config.isBetterThan(best_value->config,
-                                                 &target_density)) {
-        // Claim the previous value so that it is not included in the base.
-        (*config_claimed_map)[best_value] = true;
-        best_value = this_value;
-      } else {
-        // Claim this value so that it is not included in the base.
+      if (best_values.find(this_value) == best_values.end()) {
         (*config_claimed_map)[this_value] = true;
       }
     }
-    CHECK(best_value != nullptr);
   }
 }
 bool TableSplitter::VerifySplitConstraints(IAaptContext* context) {
@@ -263,8 +270,8 @@ void TableSplitter::SplitTable(ResourceTable* original_table) {
           }
         }
 
-        if (options_.preferred_density) {
-          MarkNonPreferredDensitiesAsClaimed(options_.preferred_density.value(),
+        if (!options_.preferred_densities.empty()) {
+          MarkNonPreferredDensitiesAsClaimed(options_.preferred_densities,
                                              density_groups,
                                              &config_claimed_map);
         }
diff --git a/tools/aapt2/split/TableSplitter.h b/tools/aapt2/split/TableSplitter.h
index 1ae327120796..6aec2576261e 100644
--- a/tools/aapt2/split/TableSplitter.h
+++ b/tools/aapt2/split/TableSplitter.h
@@ -34,14 +34,14 @@ struct SplitConstraints {
 
 struct TableSplitterOptions {
   /**
-   * The preferred density to keep in the table, stripping out all others.
+   * The preferred densities to keep in the table, stripping out all others.
+   * If empty, no stripping is done.
    */
-  Maybe<uint16_t> preferred_density;
+  std::vector<uint16_t> preferred_densities;
 
   /**
    * Configuration filter that determines which resource configuration values
-   * end up in
-   * the final table.
+   * end up in the final table.
    */
   IConfigFilter* config_filter = nullptr;
 };
diff --git a/tools/aapt2/split/TableSplitter_test.cpp b/tools/aapt2/split/TableSplitter_test.cpp
index 088dac374458..d52f4b446b8f 100644
--- a/tools/aapt2/split/TableSplitter_test.cpp
+++ b/tools/aapt2/split/TableSplitter_test.cpp
@@ -39,7 +39,7 @@ TEST(TableSplitterTest, NoSplitPreferredDensity) {
           .Build();
 
   TableSplitterOptions options;
-  options.preferred_density = ConfigDescription::DENSITY_XHIGH;
+  options.preferred_densities.push_back(ConfigDescription::DENSITY_XHIGH);
   TableSplitter splitter({}, options);
   splitter.SplitTable(table.get());
 
@@ -58,6 +58,51 @@ TEST(TableSplitterTest, NoSplitPreferredDensity) {
   EXPECT_NE(nullptr, test::GetValue<Id>(table.get(), "android:string/one"));
 }
 
+TEST(TableSplitterTest, NoSplitMultiplePreferredDensities) {
+  std::unique_ptr<ResourceTable> table =
+      test::ResourceTableBuilder()
+          .AddFileReference("android:drawable/icon",
+                            "res/drawable-mdpi/icon.png",
+                            test::ParseConfigOrDie("mdpi"))
+          .AddFileReference("android:drawable/icon",
+                            "res/drawable-hdpi/icon.png",
+                            test::ParseConfigOrDie("hdpi"))
+          .AddFileReference("android:drawable/icon",
+                            "res/drawable-xhdpi/icon.png",
+                            test::ParseConfigOrDie("xhdpi"))
+          .AddFileReference("android:drawable/icon",
+                            "res/drawable-xxhdpi/icon.png",
+                            test::ParseConfigOrDie("xxhdpi"))
+          .AddSimple("android:string/one")
+          .Build();
+
+  TableSplitterOptions options;
+  options.preferred_densities.push_back(ConfigDescription::DENSITY_LOW);
+  options.preferred_densities.push_back(ConfigDescription::DENSITY_XXXHIGH);
+  TableSplitter splitter({}, options);
+  splitter.SplitTable(table.get());
+
+  // Densities remaining:
+  // "mdpi" is the closest available density for the requested "ldpi" density.
+  EXPECT_NE(nullptr, test::GetValueForConfig<FileReference>(
+                         table.get(), "android:drawable/icon",
+                         test::ParseConfigOrDie("mdpi")));
+  // "xxhdpi" is the closest available density for the requested "xxxhdpi" density.
+  EXPECT_NE(nullptr, test::GetValueForConfig<FileReference>(
+                         table.get(), "android:drawable/icon",
+                         test::ParseConfigOrDie("xxhdpi")));
+  EXPECT_NE(nullptr, test::GetValue<Id>(table.get(), "android:string/one"));
+
+  // Removed densities:
+  EXPECT_EQ(nullptr, test::GetValueForConfig<FileReference>(
+                         table.get(), "android:drawable/icon",
+                         test::ParseConfigOrDie("hdpi")));
+  EXPECT_EQ(nullptr, test::GetValueForConfig<FileReference>(
+                         table.get(), "android:drawable/icon",
+                         test::ParseConfigOrDie("xhdpi")));
+}
+
+
 TEST(TableSplitterTest, SplitTableByDensity) {
   std::unique_ptr<ResourceTable> table =
       test::ResourceTableBuilder()
diff --git a/tools/aapt2/strip/Strip.cpp b/tools/aapt2/strip/Strip.cpp
index c34cfbf3580e..7260649c26d4 100644
--- a/tools/aapt2/strip/Strip.cpp
+++ b/tools/aapt2/strip/Strip.cpp
@@ -79,18 +79,13 @@ class StripCommand {
       context_->GetDiagnostics()->Note(DiagMessage() << "Stripping APK...");
     }
 
-    // TODO(lecesne): Add support for more than one density.
-    if (options_.target_configs.size() > 1) {
-      context_->GetDiagnostics()->Error(DiagMessage()
-                                        << "Multiple densities not supported at the moment");
-      return 1;
-    }
-
     // Stripping the APK using the TableSplitter with no splits and the target
-    // density as the preferred density. The resource table is modified in
+    // densities as the preferred densities. The resource table is modified in
     // place in the LoadedApk.
     TableSplitterOptions splitter_options;
-    splitter_options.preferred_density = options_.target_configs[0].density;
+    for (auto& config : options_.target_configs) {
+      splitter_options.preferred_densities.push_back(config.density);
+    }
     std::vector<SplitConstraints> splits;
     TableSplitter splitter(splits, splitter_options);
     splitter.SplitTable(apk->GetResourceTable());