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

10 files changed, 180 insertions, 33 deletions

diff --git a/core/java/android/content/res/AssetManager.java b/core/java/android/content/res/AssetManager.java
index 13433dc7979f..52e54af99375 100644
--- a/core/java/android/content/res/AssetManager.java
+++ b/core/java/android/content/res/AssetManager.java
@@ -43,6 +43,7 @@ import java.io.FileDescriptor;
 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 @@ public final class AssetManager implements AutoCloseable {
         }
     }
 
+    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 @@ public final class AssetManager implements AutoCloseable {
     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 ac4b7b7dc158..12e41e299e16 100644
--- a/core/java/android/content/res/Resources.java
+++ b/core/java/android/content/res/Resources.java
@@ -341,7 +341,7 @@ public class Resources {
 
     /**
      * 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 @@ public class Resources {
         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 @@ public class Resources {
             }
         }
 
-        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 @@ public class Resources {
             }
         }
 
+        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 553e11b46da5..819b01dfc58c 100644
--- a/core/java/android/content/res/ResourcesImpl.java
+++ b/core/java/android/content/res/ResourcesImpl.java
@@ -1265,16 +1265,6 @@ public class ResourcesImpl {
         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 @@ public class ResourcesImpl {
         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 @@ public class ResourcesImpl {
          * {@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 ce847e8f70c5..adff37d82fa9 100644
--- a/core/jni/android_util_AssetManager.cpp
+++ b/core/jni/android_util_AssetManager.cpp
@@ -1264,6 +1264,38 @@ static void NativeThemeApplyStyle(JNIEnv* env, jclass /*clazz*/, jlong ptr, jlon
   // 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 NativeThemeCopy(JNIEnv* env, jclass /*clazz*/, jlong dst_asset_manag
   }
 }
 
-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 @@ static const JNINativeMethod gAssetManagerMethods[] = {
     {"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 4fdae32350cf..0cde3d1242c8 100644
--- a/libs/androidfw/AssetManager2.cpp
+++ b/libs/androidfw/AssetManager2.cpp
@@ -1417,6 +1417,18 @@ base::expected<std::monostate, NullOrIOError> Theme::ApplyStyle(uint32_t resid,
   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 ae07e4bd0282..7d01395bbbbc 100644
--- a/libs/androidfw/include/androidfw/AssetManager2.h
+++ b/libs/androidfw/include/androidfw/AssetManager2.h
@@ -508,6 +508,11 @@ class Theme {
   // 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 f658735da515..77114f273d3d 100644
--- a/libs/androidfw/tests/Theme_test.cpp
+++ b/libs/androidfw/tests/Theme_test.cpp
@@ -251,6 +251,80 @@ TEST_F(ThemeTest, CopyThemeSameAssetManager) {
   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 f11486fe924a..393a20780c79 100644
--- a/libs/androidfw/tests/data/styles/R.h
+++ b/libs/androidfw/tests/data/styles/R.h
@@ -52,6 +52,7 @@ struct R {
       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 06774a8a6005..fe46a3fa90c8 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 92e9bf90101e..91cd6546c336 100644
--- a/libs/androidfw/tests/data/styles/styles.apk
+++ b/libs/androidfw/tests/data/styles/styles.apk