10 files changed, 338 insertions, 333 deletions

diff --git a/core/java/android/content/res/AssetManager.java b/core/java/android/content/res/AssetManager.java
index d50404e94544..24c6a5a12178 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;
@@ -1178,11 +1179,14 @@ public final class AssetManager implements AutoCloseable {
 
     void releaseTheme(long themePtr) {
         synchronized (this) {
-            nativeThemeDestroy(themePtr);
             decRefsLocked(themePtr);
         }
     }
 
+    static long getThemeFreeFunction() {
+        return nativeGetThemeFreeFunction();
+    }
+
     void applyStyleToTheme(long themePtr, @StyleRes int resId, boolean force) {
         synchronized (this) {
             // Need to synchronize on AssetManager because we will be accessing
@@ -1192,6 +1196,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) {
@@ -1559,12 +1588,13 @@ public final class AssetManager implements AutoCloseable {
 
     // Theme related native methods
     private static native long nativeThemeCreate(long ptr);
-    private static native void nativeThemeDestroy(long themePtr);
+    private static native long nativeGetThemeFreeFunction();
     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..b9f93b85f0bf 100644
--- a/core/java/android/content/res/ResourcesImpl.java
+++ b/core/java/android/content/res/ResourcesImpl.java
@@ -54,6 +54,8 @@ import android.view.DisplayAdjustments;
 
 import com.android.internal.util.GrowingArrayUtils;
 
+import libcore.util.NativeAllocationRegistry;
+
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
 
@@ -1265,15 +1267,9 @@ 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;
-    }
+    private static final NativeAllocationRegistry sThemeRegistry =
+            NativeAllocationRegistry.createMalloced(ResourcesImpl.class.getClassLoader(),
+                    AssetManager.getThemeFreeFunction());
 
     public class ThemeImpl {
         /**
@@ -1282,7 +1278,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;
 
         /**
@@ -1293,6 +1289,7 @@ public class ResourcesImpl {
         /*package*/ ThemeImpl() {
             mAssets = ResourcesImpl.this.mAssets;
             mTheme = mAssets.createTheme();
+            sThemeRegistry.registerNativeAllocation(this, mTheme);
         }
 
         @Override
@@ -1404,14 +1401,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..73e7d86e8279 100644
--- a/core/jni/android_util_AssetManager.cpp
+++ b/core/jni/android_util_AssetManager.cpp
@@ -1244,10 +1244,14 @@ static jlong NativeThemeCreate(JNIEnv* /*env*/, jclass /*clazz*/, jlong ptr) {
   return reinterpret_cast<jlong>(assetmanager->NewTheme().release());
 }
 
-static void NativeThemeDestroy(JNIEnv* /*env*/, jclass /*clazz*/, jlong theme_ptr) {
+static void NativeThemeDestroy(jlong theme_ptr) {
   delete reinterpret_cast<Theme*>(theme_ptr);
 }
 
+static jlong NativeGetThemeFreeFunction(JNIEnv* /*env*/, jclass /*clazz*/) {
+  return static_cast<jlong>(reinterpret_cast<uintptr_t>(&NativeThemeDestroy));
+}
+
 static void NativeThemeApplyStyle(JNIEnv* env, jclass /*clazz*/, jlong ptr, jlong theme_ptr,
                                   jint resid, jboolean force) {
   // AssetManager is accessed via the theme, so grab an explicit lock here.
@@ -1264,6 +1268,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 +1320,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) {
@@ -1446,10 +1478,11 @@ static const JNINativeMethod gAssetManagerMethods[] = {
 
     // Theme related methods.
     {"nativeThemeCreate", "(J)J", (void*)NativeThemeCreate},
-    {"nativeThemeDestroy", "(J)V", (void*)NativeThemeDestroy},
+    {"nativeGetThemeFreeFunction", "()J", (void*)NativeGetThemeFreeFunction},
     {"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 3a0153fa3dd8..0cde3d1242c8 100644
--- a/libs/androidfw/AssetManager2.cpp
+++ b/libs/androidfw/AssetManager2.cpp
@@ -1345,7 +1345,10 @@ uint8_t AssetManager2::GetAssignedPackageId(const LoadedPackage* package) const
 }
 
 std::unique_ptr<Theme> AssetManager2::NewTheme() {
-  return std::unique_ptr<Theme>(new Theme(this));
+  constexpr size_t kInitialReserveSize = 32;
+  auto theme = std::unique_ptr<Theme>(new Theme(this));
+  theme->entries_.reserve(kInitialReserveSize);
+  return theme;
 }
 
 Theme::Theme(AssetManager2* asset_manager) : asset_manager_(asset_manager) {
@@ -1353,28 +1356,20 @@ Theme::Theme(AssetManager2* asset_manager) : asset_manager_(asset_manager) {
 
 Theme::~Theme() = default;
 
-namespace {
-
-struct ThemeEntry {
+struct Theme::Entry {
+  uint32_t attr_res_id;
   ApkAssetsCookie cookie;
   uint32_t type_spec_flags;
   Res_value value;
 };
 
-struct ThemeType {
-  int entry_count;
-  ThemeEntry entries[0];
-};
-
-constexpr size_t kTypeCount = std::numeric_limits<uint8_t>::max() + 1;
-
-}  // namespace
-
-struct Theme::Package {
-  // Each element of Type will be a dynamically sized object
-  // allocated to have the entries stored contiguously with the Type.
-  std::array<util::unique_cptr<ThemeType>, kTypeCount> types;
+namespace {
+struct ThemeEntryKeyComparer {
+  bool operator() (const Theme::Entry& entry, uint32_t attr_res_id) const noexcept {
+    return entry.attr_res_id < attr_res_id;
+  }
 };
+} // namespace
 
 base::expected<std::monostate, NullOrIOError> Theme::ApplyStyle(uint32_t resid, bool force) {
   ATRACE_NAME("Theme::ApplyStyle");
@@ -1387,116 +1382,74 @@ base::expected<std::monostate, NullOrIOError> Theme::ApplyStyle(uint32_t resid,
   // Merge the flags from this style.
   type_spec_flags_ |= (*bag)->type_spec_flags;
 
-  int last_type_idx = -1;
-  int last_package_idx = -1;
-  Package* last_package = nullptr;
-  ThemeType* last_type = nullptr;
-
-  // Iterate backwards, because each bag is sorted in ascending key ID order, meaning we will only
-  // need to perform one resize per type.
-  using reverse_bag_iterator = std::reverse_iterator<const ResolvedBag::Entry*>;
-  const auto rbegin = reverse_bag_iterator(begin(*bag));
-  for (auto it = reverse_bag_iterator(end(*bag)); it != rbegin; ++it) {
-    const uint32_t attr_resid = it->key;
+  for (auto it = begin(*bag); it != end(*bag); ++it) {
+    const uint32_t attr_res_id = it->key;
 
     // If the resource ID passed in is not a style, the key can be some other identifier that is not
     // a resource ID. We should fail fast instead of operating with strange resource IDs.
-    if (!is_valid_resid(attr_resid)) {
+    if (!is_valid_resid(attr_res_id)) {
       return base::unexpected(std::nullopt);
     }
 
-    // We don't use the 0-based index for the type so that we can avoid doing ID validation
-    // upon lookup. Instead, we keep space for the type ID 0 in our data structures. Since
-    // the construction of this type is guarded with a resource ID check, it will never be
-    // populated, and querying type ID 0 will always fail.
-    const int package_idx = get_package_id(attr_resid);
-    const int type_idx = get_type_id(attr_resid);
-    const int entry_idx = get_entry_id(attr_resid);
-
-    if (last_package_idx != package_idx) {
-      std::unique_ptr<Package>& package = packages_[package_idx];
-      if (package == nullptr) {
-        package.reset(new Package());
-      }
-      last_package_idx = package_idx;
-      last_package = package.get();
-      last_type_idx = -1;
+    // DATA_NULL_EMPTY (@empty) is a valid resource value and DATA_NULL_UNDEFINED represents
+    // an absence of a valid value.
+    bool is_undefined = it->value.dataType == Res_value::TYPE_NULL &&
+        it->value.data != Res_value::DATA_NULL_EMPTY;
+    if (!force && is_undefined) {
+      continue;
     }
 
-    if (last_type_idx != type_idx) {
-      util::unique_cptr<ThemeType>& type = last_package->types[type_idx];
-      if (type == nullptr) {
-        // Allocate enough memory to contain this entry_idx. Since we're iterating in reverse over
-        // a sorted list of attributes, this shouldn't be resized again during this method call.
-        type.reset(reinterpret_cast<ThemeType*>(
-            calloc(sizeof(ThemeType) + (entry_idx + 1) * sizeof(ThemeEntry), 1)));
-        type->entry_count = entry_idx + 1;
-      } else if (entry_idx >= type->entry_count) {
-        // Reallocate the memory to contain this entry_idx. Since we're iterating in reverse over
-        // a sorted list of attributes, this shouldn't be resized again during this method call.
-        const int new_count = entry_idx + 1;
-        type.reset(reinterpret_cast<ThemeType*>(
-            realloc(type.release(), sizeof(ThemeType) + (new_count * sizeof(ThemeEntry)))));
-
-        // Clear out the newly allocated space (which isn't zeroed).
-        memset(type->entries + type->entry_count, 0,
-               (new_count - type->entry_count) * sizeof(ThemeEntry));
-        type->entry_count = new_count;
+    Theme::Entry new_entry{attr_res_id, it->cookie, (*bag)->type_spec_flags, it->value};
+    auto entry_it = std::lower_bound(entries_.begin(), entries_.end(), attr_res_id,
+                                     ThemeEntryKeyComparer{});
+    if (entry_it != entries_.end() && entry_it->attr_res_id == attr_res_id) {
+      if (is_undefined) {
+        // DATA_NULL_UNDEFINED clears the value of the attribute in the theme only when `force` is
+        /// true.
+        entries_.erase(entry_it);
+      } else if (force) {
+        *entry_it = new_entry;
       }
-      last_type_idx = type_idx;
-      last_type = type.get();
-    }
-
-    ThemeEntry& entry = last_type->entries[entry_idx];
-    if (force || (entry.value.dataType == Res_value::TYPE_NULL &&
-                  entry.value.data != Res_value::DATA_NULL_EMPTY)) {
-      entry.cookie = it->cookie;
-      entry.type_spec_flags |= (*bag)->type_spec_flags;
-      entry.value = it->value;
+    } else {
+      entries_.insert(entry_it, new_entry);
     }
   }
   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 {
 
-  int cnt = 20;
+  constexpr const uint32_t kMaxIterations = 20;
   uint32_t type_spec_flags = 0u;
-  do {
-    const int package_idx = get_package_id(resid);
-    const Package* package = packages_[package_idx].get();
-    if (package != nullptr) {
-      // The themes are constructed with a 1-based type ID, so no need to decrement here.
-      const int type_idx = get_type_id(resid);
-      const ThemeType* type = package->types[type_idx].get();
-      if (type != nullptr) {
-        const int entry_idx = get_entry_id(resid);
-        if (entry_idx < type->entry_count) {
-          const ThemeEntry& entry = type->entries[entry_idx];
-          type_spec_flags |= entry.type_spec_flags;
-
-          if (entry.value.dataType == Res_value::TYPE_ATTRIBUTE) {
-            if (cnt > 0) {
-              cnt--;
-              resid = entry.value.data;
-              continue;
-            }
-            return std::nullopt;
-          }
-
-          // @null is different than @empty.
-          if (entry.value.dataType == Res_value::TYPE_NULL &&
-              entry.value.data != Res_value::DATA_NULL_EMPTY) {
-            return std::nullopt;
-          }
+  for (uint32_t i = 0; i <= kMaxIterations; i++) {
+    auto entry_it = std::lower_bound(entries_.begin(), entries_.end(), resid,
+                                     ThemeEntryKeyComparer{});
+    if (entry_it == entries_.end() || entry_it->attr_res_id != resid) {
+      return std::nullopt;
+    }
 
-          return AssetManager2::SelectedValue(entry.value.dataType, entry.value.data, entry.cookie,
-                                              type_spec_flags, 0U /* resid */, {} /* config */);
-        }
-      }
+    type_spec_flags |= entry_it->type_spec_flags;
+    if (entry_it->value.dataType == Res_value::TYPE_ATTRIBUTE) {
+      resid = entry_it->value.data;
+      continue;
     }
-    break;
-  } while (true);
+
+    return AssetManager2::SelectedValue(entry_it->value.dataType, entry_it->value.data,
+                                        entry_it->cookie, type_spec_flags, 0U /* resid */,
+                                        {} /* config */);
+  }
   return std::nullopt;
 }
 
@@ -1520,56 +1473,25 @@ base::expected<std::monostate, NullOrIOError> Theme::ResolveAttributeReference(
 }
 
 void Theme::Clear() {
-  type_spec_flags_ = 0u;
-  for (std::unique_ptr<Package>& package : packages_) {
-    package.reset();
-  }
+  entries_.clear();
 }
 
-base::expected<std::monostate, IOError> Theme::SetTo(const Theme& o) {
-  if (this == &o) {
+base::expected<std::monostate, IOError> Theme::SetTo(const Theme& source) {
+  if (this == &source) {
     return {};
   }
 
-  type_spec_flags_ = o.type_spec_flags_;
-
-  if (asset_manager_ == o.asset_manager_) {
-    // The theme comes from the same asset manager so all theme data can be copied exactly
-    for (size_t p = 0; p < packages_.size(); p++) {
-      const Package *package = o.packages_[p].get();
-      if (package == nullptr) {
-        // The other theme doesn't have this package, clear ours.
-        packages_[p].reset();
-        continue;
-      }
-
-      if (packages_[p] == nullptr) {
-        // The other theme has this package, but we don't. Make one.
-        packages_[p].reset(new Package());
-      }
-
-      for (size_t t = 0; t < package->types.size(); t++) {
-        const ThemeType *type = package->types[t].get();
-        if (type == nullptr) {
-          // The other theme doesn't have this type, clear ours.
-          packages_[p]->types[t].reset();
-          continue;
-        }
+  type_spec_flags_ = source.type_spec_flags_;
 
-        // Create a new type and update it to theirs.
-        const size_t type_alloc_size = sizeof(ThemeType) + (type->entry_count * sizeof(ThemeEntry));
-        void *copied_data = malloc(type_alloc_size);
-        memcpy(copied_data, type, type_alloc_size);
-        packages_[p]->types[t].reset(reinterpret_cast<ThemeType *>(copied_data));
-      }
-    }
+  if (asset_manager_ == source.asset_manager_) {
+    entries_ = source.entries_;
   } else {
     std::map<ApkAssetsCookie, ApkAssetsCookie> src_to_dest_asset_cookies;
     typedef std::map<int, int> SourceToDestinationRuntimePackageMap;
     std::map<ApkAssetsCookie, SourceToDestinationRuntimePackageMap> src_asset_cookie_id_map;
 
     // Determine which ApkAssets are loaded in both theme AssetManagers.
-    const auto src_assets = o.asset_manager_->GetApkAssets();
+    const auto src_assets = source.asset_manager_->GetApkAssets();
     for (size_t i = 0; i < src_assets.size(); i++) {
       const ApkAssets* src_asset = src_assets[i];
 
@@ -1587,7 +1509,8 @@ base::expected<std::monostate, IOError> Theme::SetTo(const Theme& o) {
         // asset in th destination AssetManager.
         SourceToDestinationRuntimePackageMap package_map;
         for (const auto& loaded_package : src_asset->GetLoadedArsc()->GetPackages()) {
-          const int src_package_id = o.asset_manager_->GetAssignedPackageId(loaded_package.get());
+          const int src_package_id = source.asset_manager_->GetAssignedPackageId(
+              loaded_package.get());
           const int dest_package_id = asset_manager_->GetAssignedPackageId(loaded_package.get());
           package_map[src_package_id] = dest_package_id;
         }
@@ -1599,130 +1522,88 @@ base::expected<std::monostate, IOError> Theme::SetTo(const Theme& o) {
     }
 
     // Reset the data in the destination theme.
-    for (size_t p = 0; p < packages_.size(); p++) {
-      if (packages_[p] != nullptr) {
-        packages_[p].reset();
-      }
-    }
-
-    for (size_t p = 0; p < packages_.size(); p++) {
-      const Package *package = o.packages_[p].get();
-      if (package == nullptr) {
-        continue;
-      }
-
-      for (size_t t = 0; t < package->types.size(); t++) {
-        const ThemeType *type = package->types[t].get();
-        if (type == nullptr) {
+    entries_.clear();
+
+    for (const auto& entry : source.entries_) {
+      bool is_reference = (entry.value.dataType == Res_value::TYPE_ATTRIBUTE
+                           || entry.value.dataType == Res_value::TYPE_REFERENCE
+                           || entry.value.dataType == Res_value::TYPE_DYNAMIC_ATTRIBUTE
+                           || entry.value.dataType == Res_value::TYPE_DYNAMIC_REFERENCE)
+                          && entry.value.data != 0x0;
+
+      // If the attribute value represents an attribute or reference, the package id of the
+      // value needs to be rewritten to the package id of the value in the destination.
+      uint32_t attribute_data = entry.value.data;
+      if (is_reference) {
+        // Determine the package id of the reference in the destination AssetManager.
+        auto value_package_map = src_asset_cookie_id_map.find(entry.cookie);
+        if (value_package_map == src_asset_cookie_id_map.end()) {
           continue;
         }
 
-        for (size_t e = 0; e < type->entry_count; e++) {
-          const ThemeEntry &entry = type->entries[e];
-          if (entry.value.dataType == Res_value::TYPE_NULL &&
-              entry.value.data != Res_value::DATA_NULL_EMPTY) {
-            continue;
-          }
-
-          bool is_reference = (entry.value.dataType == Res_value::TYPE_ATTRIBUTE
-                               || entry.value.dataType == Res_value::TYPE_REFERENCE
-                               || entry.value.dataType == Res_value::TYPE_DYNAMIC_ATTRIBUTE
-                               || entry.value.dataType == Res_value::TYPE_DYNAMIC_REFERENCE)
-                              && entry.value.data != 0x0;
-
-          // If the attribute value represents an attribute or reference, the package id of the
-          // value needs to be rewritten to the package id of the value in the destination.
-          uint32_t attribute_data = entry.value.data;
-          if (is_reference) {
-            // Determine the package id of the reference in the destination AssetManager.
-            auto value_package_map = src_asset_cookie_id_map.find(entry.cookie);
-            if (value_package_map == src_asset_cookie_id_map.end()) {
-              continue;
-            }
-
-            auto value_dest_package = value_package_map->second.find(
-                get_package_id(entry.value.data));
-            if (value_dest_package == value_package_map->second.end()) {
-              continue;
-            }
-
-            attribute_data = fix_package_id(entry.value.data, value_dest_package->second);
-          }
-
-          // Find the cookie of the value in the destination. If the source apk is not loaded in the
-          // destination, only copy resources that do not reference resources in the source.
-          ApkAssetsCookie data_dest_cookie;
-          auto value_dest_cookie = src_to_dest_asset_cookies.find(entry.cookie);
-          if (value_dest_cookie != src_to_dest_asset_cookies.end()) {
-            data_dest_cookie = value_dest_cookie->second;
-          } else {
-            if (is_reference || entry.value.dataType == Res_value::TYPE_STRING) {
-              continue;
-            } else {
-              data_dest_cookie = 0x0;
-            }
-          }
+        auto value_dest_package = value_package_map->second.find(
+            get_package_id(entry.value.data));
+        if (value_dest_package == value_package_map->second.end()) {
+          continue;
+        }
 
-          // The package id of the attribute needs to be rewritten to the package id of the
-          // attribute in the destination.
-          int attribute_dest_package_id = p;
-          if (attribute_dest_package_id != 0x01) {
-            // Find the cookie of the attribute resource id in the source AssetManager
-            base::expected<FindEntryResult, NullOrIOError> attribute_entry_result =
-                o.asset_manager_->FindEntry(make_resid(p, t, e), 0 /* density_override */ ,
-                                            true /* stop_at_first_match */,
-                                            true /* ignore_configuration */);
-            if (UNLIKELY(IsIOError(attribute_entry_result))) {
-              return base::unexpected(GetIOError(attribute_entry_result.error()));
-            }
-            if (!attribute_entry_result.has_value()) {
-              continue;
-            }
-
-            // Determine the package id of the attribute in the destination AssetManager.
-            auto attribute_package_map = src_asset_cookie_id_map.find(
-                attribute_entry_result->cookie);
-            if (attribute_package_map == src_asset_cookie_id_map.end()) {
-              continue;
-            }
-            auto attribute_dest_package = attribute_package_map->second.find(
-                attribute_dest_package_id);
-            if (attribute_dest_package == attribute_package_map->second.end()) {
-              continue;
-            }
-            attribute_dest_package_id = attribute_dest_package->second;
-          }
+        attribute_data = fix_package_id(entry.value.data, value_dest_package->second);
+      }
 
-          // Lazily instantiate the destination package.
-          std::unique_ptr<Package>& dest_package = packages_[attribute_dest_package_id];
-          if (dest_package == nullptr) {
-            dest_package.reset(new Package());
-          }
+      // Find the cookie of the value in the destination. If the source apk is not loaded in the
+      // destination, only copy resources that do not reference resources in the source.
+      ApkAssetsCookie data_dest_cookie;
+      auto value_dest_cookie = src_to_dest_asset_cookies.find(entry.cookie);
+      if (value_dest_cookie != src_to_dest_asset_cookies.end()) {
+        data_dest_cookie = value_dest_cookie->second;
+      } else {
+        if (is_reference || entry.value.dataType == Res_value::TYPE_STRING) {
+          continue;
+        } else {
+          data_dest_cookie = 0x0;
+        }
+      }
 
-          // Lazily instantiate and resize the destination type.
-          util::unique_cptr<ThemeType>& dest_type = dest_package->types[t];
-          if (dest_type == nullptr || dest_type->entry_count < type->entry_count) {
-            const size_t type_alloc_size = sizeof(ThemeType)
-                + (type->entry_count * sizeof(ThemeEntry));
-            void* dest_data = malloc(type_alloc_size);
-            memset(dest_data, 0, type->entry_count * sizeof(ThemeEntry));
-
-            // Copy the existing destination type values if the type is resized.
-            if (dest_type != nullptr) {
-              memcpy(dest_data, type, sizeof(ThemeType)
-                                      + (dest_type->entry_count * sizeof(ThemeEntry)));
-            }
-
-            dest_type.reset(reinterpret_cast<ThemeType *>(dest_data));
-            dest_type->entry_count = type->entry_count;
-          }
+      // The package id of the attribute needs to be rewritten to the package id of the
+      // attribute in the destination.
+      int attribute_dest_package_id = get_package_id(entry.attr_res_id);
+      if (attribute_dest_package_id != 0x01) {
+        // Find the cookie of the attribute resource id in the source AssetManager
+        base::expected<FindEntryResult, NullOrIOError> attribute_entry_result =
+            source.asset_manager_->FindEntry(entry.attr_res_id, 0 /* density_override */ ,
+                                             true /* stop_at_first_match */,
+                                             true /* ignore_configuration */);
+        if (UNLIKELY(IsIOError(attribute_entry_result))) {
+          return base::unexpected(GetIOError(attribute_entry_result.error()));
+        }
+        if (!attribute_entry_result.has_value()) {
+          continue;
+        }
 
-          dest_type->entries[e].cookie = data_dest_cookie;
-          dest_type->entries[e].value.dataType = entry.value.dataType;
-          dest_type->entries[e].value.data = attribute_data;
-          dest_type->entries[e].type_spec_flags = entry.type_spec_flags;
+        // Determine the package id of the attribute in the destination AssetManager.
+        auto attribute_package_map = src_asset_cookie_id_map.find(
+            attribute_entry_result->cookie);
+        if (attribute_package_map == src_asset_cookie_id_map.end()) {
+          continue;
+        }
+        auto attribute_dest_package = attribute_package_map->second.find(
+            attribute_dest_package_id);
+        if (attribute_dest_package == attribute_package_map->second.end()) {
+          continue;
         }
+        attribute_dest_package_id = attribute_dest_package->second;
       }
+
+      auto dest_attr_id = make_resid(attribute_dest_package_id, get_type_id(entry.attr_res_id),
+                                     get_entry_id(entry.attr_res_id));
+      Theme::Entry new_entry{dest_attr_id, data_dest_cookie, entry.type_spec_flags,
+                                            Res_value{.dataType = entry.value.dataType,
+                                                      .data = attribute_data}};
+
+      // Since the entries were cleared, the attribute resource id has yet been mapped to any value.
+      auto entry_it = std::lower_bound(entries_.begin(), entries_.end(), dest_attr_id,
+                                       ThemeEntryKeyComparer{});
+      entries_.insert(entry_it, new_entry);
     }
   }
   return {};
@@ -1730,31 +1611,10 @@ base::expected<std::monostate, IOError> Theme::SetTo(const Theme& o) {
 
 void Theme::Dump() const {
   LOG(INFO) << base::StringPrintf("Theme(this=%p, AssetManager2=%p)", this, asset_manager_);
-
-  for (int p = 0; p < packages_.size(); p++) {
-    auto& package = packages_[p];
-    if (package == nullptr) {
-      continue;
-    }
-
-    for (int t = 0; t < package->types.size(); t++) {
-      auto& type = package->types[t];
-      if (type == nullptr) {
-        continue;
-      }
-
-      for (int e = 0; e < type->entry_count; e++) {
-        auto& entry = type->entries[e];
-        if (entry.value.dataType == Res_value::TYPE_NULL &&
-            entry.value.data != Res_value::DATA_NULL_EMPTY) {
-          continue;
-        }
-
-        LOG(INFO) << base::StringPrintf("  entry(0x%08x)=(0x%08x) type=(0x%02x), cookie(%d)",
-                                        make_resid(p, t, e), entry.value.data,
-                                        entry.value.dataType, entry.cookie);
-      }
-    }
+  for (auto& entry : entries_) {
+    LOG(INFO) << base::StringPrintf("  entry(0x%08x)=(0x%08x) type=(0x%02x), cookie(%d)",
+                                    entry.attr_res_id, entry.value.data, entry.value.dataType,
+                                    entry.cookie);
   }
 }
 
diff --git a/libs/androidfw/include/androidfw/AssetManager2.h b/libs/androidfw/include/androidfw/AssetManager2.h
index df3abf63323a..7d01395bbbbc 100644
--- a/libs/androidfw/include/androidfw/AssetManager2.h
+++ b/libs/androidfw/include/androidfw/AssetManager2.h
@@ -508,13 +508,18 @@ class Theme {
   // data failed.
   base::expected<std::monostate, NullOrIOError> ApplyStyle(uint32_t resid, bool force = false);
 
-  // Sets this Theme to be a copy of `other` if `other` has the same AssetManager as this Theme.
+  // 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 `other` does not have the same AssetManager as this theme, only attributes from ApkAssets
+  // If `source` does not have the same AssetManager as this theme, only attributes from ApkAssets
   // loaded into both AssetManagers will be copied to this theme.
   //
   // Returns an I/O error if reading resource data failed.
-  base::expected<std::monostate, IOError> SetTo(const Theme& other);
+  base::expected<std::monostate, IOError> SetTo(const Theme& source);
 
   void Clear();
 
@@ -546,20 +551,16 @@ class Theme {
 
   void Dump() const;
 
+  struct Entry;
  private:
   DISALLOW_COPY_AND_ASSIGN(Theme);
 
-  // Called by AssetManager2.
   explicit Theme(AssetManager2* asset_manager);
 
-  AssetManager2* asset_manager_;
+  AssetManager2* asset_manager_ = nullptr;
   uint32_t type_spec_flags_ = 0u;
 
-  // Defined in the cpp.
-  struct Package;
-
-  constexpr static size_t kPackageCount = std::numeric_limits<uint8_t>::max() + 1;
-  std::array<std::unique_ptr<Package>, kPackageCount> packages_;
+  std::vector<Entry> entries_;
 };
 
 inline const ResolvedBag::Entry* begin(const ResolvedBag* bag) {
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