Rebase ThemeImpl rather than reallocate memory

Memory churn is high when swapping the ResourcesImpl of a Resources
object. Each time Resources#setImpl is invoked, all themes based on
that Resources object are assigned new ThemeImpl objects that are
created using the new ResourcesImpl.

ThemeImpls can only belong to one Theme object, so the old
implementation is discarded and the theme takes ownership of the new
ThemeImp.

This creates performance problems when framework overlays are toggled.
Toggling overlays targeting the framework causes all themes across all
processes to recreate and reallocate all of their themes. By rebasing
the ThemeImpl on the new ResourcesImpl without deallocating the native
theme memory, we reduce churn and produce less garbage that needs to
be garbage collected.

Bug: 141198925
Test: atest libandroidfw_tests
Test: atest ResourcesPerfWorkloads
Change-Id: I03fb31ee09c9cfdbd3c41bcf0b605607dab54ed7
diff --git a/libs/androidfw/AssetManager2.cpp b/libs/androidfw/AssetManager2.cpp
index 4fdae32..0cde3d1 100644
--- a/libs/androidfw/AssetManager2.cpp
+++ b/libs/androidfw/AssetManager2.cpp
@@ -1417,6 +1417,18 @@
   return {};
 }
 
+void Theme::Rebase(AssetManager2* am, const uint32_t* style_ids, const uint8_t* force,
+                   size_t style_count) {
+  ATRACE_NAME("Theme::Rebase");
+  // Reset the entries without changing the vector capacity to prevent reallocations during
+  // ApplyStyle.
+  entries_.clear();
+  asset_manager_ = am;
+  for (size_t i = 0; i < style_count; i++) {
+    ApplyStyle(style_ids[i], force[i]);
+  }
+}
+
 std::optional<AssetManager2::SelectedValue> Theme::GetAttribute(uint32_t resid) const {
 
   constexpr const uint32_t kMaxIterations = 20;
diff --git a/libs/androidfw/include/androidfw/AssetManager2.h b/libs/androidfw/include/androidfw/AssetManager2.h
index ae07e4b..7d01395 100644
--- a/libs/androidfw/include/androidfw/AssetManager2.h
+++ b/libs/androidfw/include/androidfw/AssetManager2.h
@@ -508,6 +508,11 @@
   // data failed.
   base::expected<std::monostate, NullOrIOError> ApplyStyle(uint32_t resid, bool force = false);
 
+  // Clears the existing theme, sets the new asset manager to use for this theme, and applies the
+  // styles in `style_ids` through repeated invocations of `ApplyStyle`.
+  void Rebase(AssetManager2* am, const uint32_t* style_ids, const uint8_t* force,
+              size_t style_count);
+
   // Sets this Theme to be a copy of `source` if `source` has the same AssetManager as this Theme.
   //
   // If `source` does not have the same AssetManager as this theme, only attributes from ApkAssets
diff --git a/libs/androidfw/tests/Theme_test.cpp b/libs/androidfw/tests/Theme_test.cpp
index f658735..77114f2 100644
--- a/libs/androidfw/tests/Theme_test.cpp
+++ b/libs/androidfw/tests/Theme_test.cpp
@@ -251,6 +251,80 @@
   EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC), value->flags);
 }
 
+TEST_F(ThemeTest, ThemeRebase) {
+  AssetManager2 am;
+  am.SetApkAssets({style_assets_.get()});
+
+  AssetManager2 am_night;
+  am_night.SetApkAssets({style_assets_.get()});
+
+  ResTable_config night{};
+  night.uiMode = ResTable_config::UI_MODE_NIGHT_YES;
+  night.version = 8u;
+  am_night.SetConfiguration(night);
+
+  auto theme = am.NewTheme();
+  {
+    const uint32_t styles[] = {app::R::style::StyleOne, app::R::style::StyleDayNight};
+    const uint8_t force[] = {true, true};
+    theme->Rebase(&am, styles, force, arraysize(styles));
+  }
+
+  // attr_one in StyleDayNight force overrides StyleOne. attr_one is defined in the StyleOne.
+  auto value = theme->GetAttribute(app::R::attr::attr_one);
+  ASSERT_TRUE(value);
+  EXPECT_EQ(10u, value->data);
+  EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC | ResTable_config::CONFIG_UI_MODE |
+            ResTable_config::CONFIG_VERSION), value->flags);
+
+  // attr_two is defined in the StyleOne.
+  value = theme->GetAttribute(app::R::attr::attr_two);
+  ASSERT_TRUE(value);
+  EXPECT_EQ(Res_value::TYPE_INT_DEC, value->type);
+  EXPECT_EQ(2u, value->data);
+  EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC), value->flags);
+
+  {
+    const uint32_t styles[] = {app::R::style::StyleOne, app::R::style::StyleDayNight};
+    const uint8_t force[] = {false, false};
+    theme->Rebase(&am, styles, force, arraysize(styles));
+  }
+
+  // attr_one in StyleDayNight does not override StyleOne because `force` is false.
+  value = theme->GetAttribute(app::R::attr::attr_one);
+  ASSERT_TRUE(value);
+  EXPECT_EQ(1u, value->data);
+  EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC), value->flags);
+
+  // attr_two is defined in the StyleOne.
+  value = theme->GetAttribute(app::R::attr::attr_two);
+  ASSERT_TRUE(value);
+  EXPECT_EQ(Res_value::TYPE_INT_DEC, value->type);
+  EXPECT_EQ(2u, value->data);
+  EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC), value->flags);
+
+  {
+    const uint32_t styles[] = {app::R::style::StyleOne, app::R::style::StyleDayNight};
+    const uint8_t force[] = {false, true};
+    theme->Rebase(&am_night, styles, force, arraysize(styles));
+  }
+
+  // attr_one is defined in the StyleDayNight.
+  value = theme->GetAttribute(app::R::attr::attr_one);
+  ASSERT_TRUE(value);
+  EXPECT_EQ(Res_value::TYPE_INT_DEC, value->type);
+  EXPECT_EQ(100u, value->data);
+  EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC | ResTable_config::CONFIG_UI_MODE |
+            ResTable_config::CONFIG_VERSION), value->flags);
+
+  // attr_two is now not here.
+  value = theme->GetAttribute(app::R::attr::attr_two);
+  ASSERT_TRUE(value);
+  EXPECT_EQ(Res_value::TYPE_INT_DEC, value->type);
+  EXPECT_EQ(2u, value->data);
+  EXPECT_EQ(static_cast<uint32_t>(ResTable_typeSpec::SPEC_PUBLIC), value->flags);
+}
+
 TEST_F(ThemeTest, OnlyCopySameAssetsThemeWhenAssetManagersDiffer) {
   AssetManager2 assetmanager_dst;
   assetmanager_dst.SetApkAssets({system_assets_.get(), lib_one_assets_.get(), style_assets_.get(),
diff --git a/libs/androidfw/tests/data/styles/R.h b/libs/androidfw/tests/data/styles/R.h
index f11486f..393a207 100644
--- a/libs/androidfw/tests/data/styles/R.h
+++ b/libs/androidfw/tests/data/styles/R.h
@@ -52,6 +52,7 @@
       StyleFive = 0x7f020004u,
       StyleSix = 0x7f020005u,
       StyleSeven = 0x7f020006u,
+      StyleDayNight = 0x7f020007u,
     };
   };
 };
diff --git a/libs/androidfw/tests/data/styles/res/values/styles.xml b/libs/androidfw/tests/data/styles/res/values/styles.xml
index 06774a8..fe46a3fa 100644
--- a/libs/androidfw/tests/data/styles/res/values/styles.xml
+++ b/libs/androidfw/tests/data/styles/res/values/styles.xml
@@ -89,4 +89,9 @@
         <item name="android:theme">@empty</item>
     </style>
 
+    <public type="style" name="StyleDayNight" id="0x7f020007" />
+    <style name="StyleDayNight">
+        <item name="attr_one">10</item>
+    </style>
+
 </resources>
diff --git a/libs/androidfw/tests/data/styles/styles.apk b/libs/androidfw/tests/data/styles/styles.apk
index 92e9bf9..91cd654 100644
--- a/libs/androidfw/tests/data/styles/styles.apk
+++ b/libs/androidfw/tests/data/styles/styles.apk
Binary files differ