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/core/java/android/content/res/AssetManager.java b/core/java/android/content/res/AssetManager.java
index 13433dc..52e54af 100644
--- a/core/java/android/content/res/AssetManager.java
+++ b/core/java/android/content/res/AssetManager.java
@@ -43,6 +43,7 @@
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.ref.Reference;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -1187,6 +1188,31 @@
         }
     }
 
+    AssetManager rebaseTheme(long themePtr, @NonNull AssetManager newAssetManager,
+            @StyleRes int[] styleIds, @StyleRes boolean[] force, int count) {
+        // Exchange ownership of the theme with the new asset manager.
+        if (this != newAssetManager) {
+            synchronized (this) {
+                ensureValidLocked();
+                decRefsLocked(themePtr);
+            }
+            synchronized (newAssetManager) {
+                newAssetManager.ensureValidLocked();
+                newAssetManager.incRefsLocked(themePtr);
+            }
+        }
+
+        try {
+            synchronized (newAssetManager) {
+                newAssetManager.ensureValidLocked();
+                nativeThemeRebase(newAssetManager.mObject, themePtr, styleIds, force, count);
+            }
+        } finally {
+            Reference.reachabilityFence(newAssetManager);
+        }
+        return newAssetManager;
+    }
+
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
     void setThemeTo(long dstThemePtr, @NonNull AssetManager srcAssetManager, long srcThemePtr) {
         synchronized (this) {
@@ -1557,9 +1583,10 @@
     private static native void nativeThemeDestroy(long themePtr);
     private static native void nativeThemeApplyStyle(long ptr, long themePtr, @StyleRes int resId,
             boolean force);
+    private static native void nativeThemeRebase(long ptr, long themePtr, @NonNull int[] styleIds,
+            @NonNull boolean[] force, int styleSize);
     private static native void nativeThemeCopy(long dstAssetManagerPtr, long dstThemePtr,
             long srcAssetManagerPtr, long srcThemePtr);
-    static native void nativeThemeClear(long themePtr);
     private static native int nativeThemeGetAttributeValue(long ptr, long themePtr,
             @AttrRes int resId, @NonNull TypedValue outValue, boolean resolve);
     private static native void nativeThemeDump(long ptr, long themePtr, int priority, String tag,
diff --git a/core/java/android/content/res/Resources.java b/core/java/android/content/res/Resources.java
index ac4b7b7..12e41e2 100644
--- a/core/java/android/content/res/Resources.java
+++ b/core/java/android/content/res/Resources.java
@@ -341,7 +341,7 @@
 
     /**
      * Set the underlying implementation (containing all the resources and caches)
-     * and updates all Theme references to new implementations as well.
+     * and updates all Theme implementations as well.
      * @hide
      */
     @UnsupportedAppUsage
@@ -353,14 +353,14 @@
         mBaseApkAssetsSize = ArrayUtils.size(impl.getAssets().getApkAssets());
         mResourcesImpl = impl;
 
-        // Create new ThemeImpls that are identical to the ones we have.
+        // Rebase the ThemeImpls using the new ResourcesImpl.
         synchronized (mThemeRefs) {
             final int count = mThemeRefs.size();
             for (int i = 0; i < count; i++) {
                 WeakReference<Theme> weakThemeRef = mThemeRefs.get(i);
                 Theme theme = weakThemeRef != null ? weakThemeRef.get() : null;
                 if (theme != null) {
-                    theme.setNewResourcesImpl(mResourcesImpl);
+                    theme.rebase(mResourcesImpl);
                 }
             }
         }
@@ -1515,12 +1515,6 @@
             }
         }
 
-        void setNewResourcesImpl(ResourcesImpl resImpl) {
-            synchronized (mLock) {
-                mThemeImpl = resImpl.newThemeImpl(mThemeImpl.getKey());
-            }
-        }
-
         /**
          * Place new attribute values into the theme.  The style resource
          * specified by <var>resid</var> will be retrieved from this Theme's
@@ -1847,6 +1841,12 @@
             }
         }
 
+        void rebase(ResourcesImpl resImpl) {
+            synchronized (mLock) {
+                mThemeImpl.rebase(resImpl.mAssets);
+            }
+        }
+
         /**
          * Returns the resource ID for the style specified using {@code style="..."} in the
          * {@link AttributeSet}'s backing XML element or {@link Resources#ID_NULL} otherwise if not
diff --git a/core/java/android/content/res/ResourcesImpl.java b/core/java/android/content/res/ResourcesImpl.java
index 553e11b..819b01d 100644
--- a/core/java/android/content/res/ResourcesImpl.java
+++ b/core/java/android/content/res/ResourcesImpl.java
@@ -1265,16 +1265,6 @@
         return new ThemeImpl();
     }
 
-    /**
-     * Creates a new ThemeImpl which is already set to the given Resources.ThemeKey.
-     */
-    ThemeImpl newThemeImpl(Resources.ThemeKey key) {
-        ThemeImpl impl = new ThemeImpl();
-        impl.mKey.setTo(key);
-        impl.rebase();
-        return impl;
-    }
-
     public class ThemeImpl {
         /**
          * Unique key for the series of styles applied to this theme.
@@ -1282,7 +1272,7 @@
         private final Resources.ThemeKey mKey = new Resources.ThemeKey();
 
         @SuppressWarnings("hiding")
-        private final AssetManager mAssets;
+        private AssetManager mAssets;
         private final long mTheme;
 
         /**
@@ -1404,14 +1394,18 @@
          * {@link #applyStyle(int, boolean)}.
          */
         void rebase() {
-            AssetManager.nativeThemeClear(mTheme);
+            rebase(mAssets);
+        }
 
-            // Reapply the same styles in the same order.
-            for (int i = 0; i < mKey.mCount; i++) {
-                final int resId = mKey.mResId[i];
-                final boolean force = mKey.mForce[i];
-                mAssets.applyStyleToTheme(mTheme, resId, force);
-            }
+        /**
+         * Rebases the theme against the {@code newAssets} by re-applying the styles passed to
+         * {@link #applyStyle(int, boolean)}.
+         *
+         * The theme will use {@code newAssets} for all future invocations of
+         * {@link #applyStyle(int, boolean)}.
+         */
+        void rebase(AssetManager newAssets) {
+            mAssets = mAssets.rebaseTheme(mTheme, newAssets, mKey.mResId, mKey.mForce, mKey.mCount);
         }
 
         /**
diff --git a/core/jni/android_util_AssetManager.cpp b/core/jni/android_util_AssetManager.cpp
index ce847e8..adff37d 100644
--- a/core/jni/android_util_AssetManager.cpp
+++ b/core/jni/android_util_AssetManager.cpp
@@ -1264,6 +1264,38 @@
   // jniThrowException(env, "java/lang/IllegalArgumentException", error_msg.c_str());
 }
 
+static void NativeThemeRebase(JNIEnv* env, jclass /*clazz*/, jlong ptr, jlong theme_ptr,
+                              jintArray style_ids, jbooleanArray force,
+                              jint style_count) {
+  // Lock both the original asset manager of the theme and the new asset manager to be used for the
+  // theme.
+  ScopedLock<AssetManager2> assetmanager(AssetManagerFromLong(ptr));
+
+  uint32_t* style_id_args = nullptr;
+  if (style_ids != nullptr) {
+    CHECK(style_count <= env->GetArrayLength(style_ids));
+    style_id_args = reinterpret_cast<uint32_t*>(env->GetPrimitiveArrayCritical(style_ids, nullptr));
+    if (style_id_args == nullptr) {
+      return;
+    }
+  }
+
+  jboolean* force_args = nullptr;
+  if (force != nullptr) {
+    CHECK(style_count <= env->GetArrayLength(force));
+    force_args = reinterpret_cast<jboolean*>(env->GetPrimitiveArrayCritical(force, nullptr));
+    if (force_args == nullptr) {
+      env->ReleasePrimitiveArrayCritical(style_ids, style_id_args, JNI_ABORT);
+      return;
+    }
+  }
+
+  auto theme = reinterpret_cast<Theme*>(theme_ptr);
+  theme->Rebase(&(*assetmanager), style_id_args, force_args, static_cast<size_t>(style_count));
+  env->ReleasePrimitiveArrayCritical(style_ids, style_id_args, JNI_ABORT);
+  env->ReleasePrimitiveArrayCritical(force, force_args, JNI_ABORT);
+}
+
 static void NativeThemeCopy(JNIEnv* env, jclass /*clazz*/, jlong dst_asset_manager_ptr,
                             jlong dst_theme_ptr, jlong src_asset_manager_ptr, jlong src_theme_ptr) {
   Theme* dst_theme = reinterpret_cast<Theme*>(dst_theme_ptr);
@@ -1284,10 +1316,6 @@
   }
 }
 
-static void NativeThemeClear(JNIEnv* /*env*/, jclass /*clazz*/, jlong theme_ptr) {
-  reinterpret_cast<Theme*>(theme_ptr)->Clear();
-}
-
 static jint NativeThemeGetAttributeValue(JNIEnv* env, jclass /*clazz*/, jlong ptr, jlong theme_ptr,
                                          jint resid, jobject typed_value,
                                          jboolean resolve_references) {
@@ -1448,8 +1476,9 @@
     {"nativeThemeCreate", "(J)J", (void*)NativeThemeCreate},
     {"nativeThemeDestroy", "(J)V", (void*)NativeThemeDestroy},
     {"nativeThemeApplyStyle", "(JJIZ)V", (void*)NativeThemeApplyStyle},
+    {"nativeThemeRebase", "(JJ[I[ZI)V", (void*)NativeThemeRebase},
+
     {"nativeThemeCopy", "(JJJJ)V", (void*)NativeThemeCopy},
-    {"nativeThemeClear", "(J)V", (void*)NativeThemeClear},
     {"nativeThemeGetAttributeValue", "(JJILandroid/util/TypedValue;Z)I",
      (void*)NativeThemeGetAttributeValue},
     {"nativeThemeDump", "(JJILjava/lang/String;Ljava/lang/String;)V", (void*)NativeThemeDump},
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