[res] Optimize Theme::ApplyStyle()

This is the heaviest operation during the configuration change
handling, as we recreate the asset manager object and rebase
all existing themes on that object as soon as any assets
change. The optimization here replaces the repeated binary
search + insert in a middle of an array with a merge of two
sorted arrays in place, achieving over 7x speedup:

Before:
  #BM_ThemeApplyStyleFramework 9722.817566754931 ns
After:
  #BM_ThemeApplyStyleFramework 1305.8760730843824 ns

It also adds a detailed explanation of the algorithm and its
assumptions to make it easier to optimize/debug it later.

Flag: NONE A performance regression fix, flag check is too slow
Bug: 345562237
Test: build, boot, atest + performance tests

Change-Id: I979e17cf4837cc8853a8d54cb4cea2a9631c3136
diff --git a/libs/androidfw/Android.bp b/libs/androidfw/Android.bp
index 260f3c1..2fff4f5 100644
--- a/libs/androidfw/Android.bp
+++ b/libs/androidfw/Android.bp
@@ -212,6 +212,7 @@
         "tests/AttributeResolution_test.cpp",
         "tests/BigBuffer_test.cpp",
         "tests/ByteBucketArray_test.cpp",
+        "tests/CombinedIterator_test.cpp",
         "tests/Config_test.cpp",
         "tests/ConfigDescription_test.cpp",
         "tests/ConfigLocale_test.cpp",
diff --git a/libs/androidfw/AssetManager2.cpp b/libs/androidfw/AssetManager2.cpp
index 46f636e..822a387 100644
--- a/libs/androidfw/AssetManager2.cpp
+++ b/libs/androidfw/AssetManager2.cpp
@@ -23,9 +23,11 @@
 #include <map>
 #include <set>
 #include <span>
+#include <utility>
 
 #include "android-base/logging.h"
 #include "android-base/stringprintf.h"
+#include "androidfw/CombinedIterator.h"
 #include "androidfw/ResourceTypes.h"
 #include "androidfw/ResourceUtils.h"
 #include "androidfw/Util.h"
@@ -1622,6 +1624,12 @@
 
 Theme::~Theme() = default;
 
+static bool IsUndefined(const Res_value& value) {
+  // DATA_NULL_EMPTY (@empty) is a valid resource value and DATA_NULL_UNDEFINED represents
+  // an absence of a valid value.
+  return value.dataType == Res_value::TYPE_NULL && value.data != Res_value::DATA_NULL_EMPTY;
+}
+
 base::expected<std::monostate, NullOrIOError> Theme::ApplyStyle(uint32_t resid, bool force) {
   ATRACE_NAME("Theme::ApplyStyle");
 
@@ -1633,39 +1641,76 @@
   // Merge the flags from this style.
   type_spec_flags_ |= (*bag)->type_spec_flags;
 
+  //
+  // This function is the most expensive part of applying an frro to the existing app resources,
+  // and needs to be as efficient as possible.
+  // The data structure we're working with is two parallel sorted arrays of keys (resource IDs)
+  // and entries (resource value + some attributes).
+  // The styles get applied in sequence, starting with an empty set of attributes. Each style
+  // contains its values for the theme attributes, and gets applied in either normal or forced way:
+  //  - normal way never overrides the existing attribute, so only unique style attributes are added
+  //  - forced way overrides anything for that attribute, and if it's undefined it removes the
+  //    previous value completely
+  //
+  // Style attributes come in a Bag data type - a sorted array of attributes with their values. This
+  // means we don't need to re-sort the attributes ever, and instead:
+  //  - for an already existing attribute just skip it or apply the forced value
+  //    - if the forced value is undefined, mark it undefined as well to get rid of it later
+  //  - for a new attribute append it to the array, forming a new sorted section of new attributes
+  //    past the end of the original ones (ignore undefined ones here)
+  //  - inplace merge two sorted sections to form a single sorted array again.
+  //  - run the last pass to remove all undefined elements
+  //
+  // Using this algorithm performs better than a repeated binary search + insert in the middle,
+  // as that keeps shifting the tail end of the arrays and wasting CPU cycles in memcpy().
+  //
+  const auto starting_size = keys_.size();
+  if (starting_size == 0) {
+    keys_.reserve((*bag)->entry_count);
+    entries_.reserve((*bag)->entry_count);
+  }
+  bool wrote_undefined = false;
   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_res_id)) {
       return base::unexpected(std::nullopt);
     }
-
-    // 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;
+    const bool is_undefined = IsUndefined(it->value);
     if (!force && is_undefined) {
       continue;
     }
-
-    const auto key_it = std::lower_bound(keys_.begin(), keys_.end(), attr_res_id);
-    const auto entry_it = entries_.begin() + (key_it - keys_.begin());
-    if (key_it != keys_.end() && *key_it == attr_res_id) {
-      if (is_undefined) {
-        // DATA_NULL_UNDEFINED clears the value of the attribute in the theme only when `force` is
-        // true.
-        keys_.erase(key_it);
-        entries_.erase(entry_it);
-      } else if (force) {
+    const auto key_it = std::lower_bound(keys_.begin(), keys_.begin() + starting_size, attr_res_id);
+    if (key_it != keys_.begin() + starting_size && *key_it == attr_res_id) {
+      const auto entry_it = entries_.begin() + (key_it - keys_.begin());
+      if (force || IsUndefined(entry_it->value)) {
         *entry_it = Entry{it->cookie, (*bag)->type_spec_flags, it->value};
+        wrote_undefined |= is_undefined;
       }
-    } else {
-      keys_.insert(key_it, attr_res_id);
-      entries_.insert(entry_it, Entry{it->cookie, (*bag)->type_spec_flags, it->value});
+    } else if (!is_undefined) {
+      keys_.emplace_back(attr_res_id);
+      entries_.emplace_back(it->cookie, (*bag)->type_spec_flags, it->value);
     }
   }
+
+  if (starting_size && keys_.size() != starting_size) {
+    std::inplace_merge(
+        CombinedIterator(keys_.begin(), entries_.begin()),
+        CombinedIterator(keys_.begin() + starting_size, entries_.begin() + starting_size),
+        CombinedIterator(keys_.end(), entries_.end()));
+  }
+  if (wrote_undefined) {
+    auto new_end = std::remove_if(CombinedIterator(keys_.begin(), entries_.begin()),
+                                  CombinedIterator(keys_.end(), entries_.end()),
+                                  [](const auto& pair) { return IsUndefined(pair.second.value); });
+    keys_.erase(new_end.it1, keys_.end());
+    entries_.erase(new_end.it2, entries_.end());
+  }
+  if (android::base::kEnableDChecks && !std::is_sorted(keys_.begin(), keys_.end())) {
+    ALOGW("Bag %u was unsorted in the apk?", unsigned(resid));
+    return base::unexpected(std::nullopt);
+  }
   return {};
 }
 
@@ -1691,6 +1736,9 @@
       return std::nullopt;
     }
     const auto entry_it = entries_.begin() + (key_it - keys_.begin());
+    if (IsUndefined(entry_it->value)) {
+      return std::nullopt;
+    }
     type_spec_flags |= entry_it->type_spec_flags;
     if (entry_it->value.dataType == Res_value::TYPE_ATTRIBUTE) {
       resid = entry_it->value.data;
diff --git a/libs/androidfw/include/androidfw/CombinedIterator.h b/libs/androidfw/include/androidfw/CombinedIterator.h
new file mode 100644
index 0000000..4ff6a7d
--- /dev/null
+++ b/libs/androidfw/include/androidfw/CombinedIterator.h
@@ -0,0 +1,176 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <compare>
+#include <iterator>
+#include <utility>
+
+namespace android {
+
+namespace detail {
+// A few useful aliases to not repeat them everywhere
+template <class It1, class It2>
+using Value = std::pair<typename std::iterator_traits<It1>::value_type,
+                        typename std::iterator_traits<It2>::value_type>;
+
+template <class It1, class It2>
+using BaseRefPair = std::pair<typename std::iterator_traits<It1>::reference,
+                              typename std::iterator_traits<It2>::reference>;
+
+template <class It1, class It2>
+struct RefPair : BaseRefPair<It1, It2> {
+  using Base = BaseRefPair<It1, It2>;
+  using Value = detail::Value<It1, It2>;
+
+  RefPair(It1 it1, It2 it2) : Base(*it1, *it2) {
+  }
+
+  RefPair& operator=(const Value& v) {
+    this->first = v.first;
+    this->second = v.second;
+    return *this;
+  }
+  operator Value() const {
+    return Value(this->first, this->second);
+  }
+  bool operator==(const RefPair& other) {
+    return this->first == other.first;
+  }
+  bool operator==(const Value& other) {
+    return this->first == other.first;
+  }
+  std::strong_ordering operator<=>(const RefPair& other) const {
+    return this->first <=> other.first;
+  }
+  std::strong_ordering operator<=>(const Value& other) const {
+    return this->first <=> other.first;
+  }
+  friend void swap(RefPair& l, RefPair& r) {
+    using std::swap;
+    swap(l.first, r.first);
+    swap(l.second, r.second);
+  }
+};
+
+template <class It1, class It2>
+struct RefPairPtr {
+  RefPair<It1, It2> value;
+
+  RefPair<It1, It2>* operator->() const {
+    return &value;
+  }
+};
+}  // namespace detail
+
+//
+//   CombinedIterator - a class to combine two iterators to process them as a single iterator to a
+// pair of values. Useful for processing a data structure of "struct of arrays", replacing
+// array of structs for cache locality.
+//
+// The value type is a pair of copies of the values of each iterator, and the reference is a
+// pair of references to the corresponding values. Comparison only compares the first element,
+// making it most useful for using on data like (vector<Key>, vector<Value>) for binary searching,
+// sorting both together and so on.
+//
+// The class is designed for handling arrays, so it requires random access iterators as an input.
+//
+
+template <class It1, class It2>
+requires std::random_access_iterator<It1> && std::random_access_iterator<It2>
+struct CombinedIterator {
+  typedef detail::Value<It1, It2> value_type;
+  typedef detail::RefPair<It1, It2> reference;
+  typedef std::ptrdiff_t difference_type;
+  typedef detail::RefPairPtr<It1, It2> pointer;
+  typedef std::random_access_iterator_tag iterator_category;
+
+  CombinedIterator(It1 it1 = {}, It2 it2 = {}) : it1(it1), it2(it2) {
+  }
+
+  bool operator<(const CombinedIterator& other) const {
+    return it1 < other.it1;
+  }
+  bool operator<=(const CombinedIterator& other) const {
+    return it1 <= other.it1;
+  }
+  bool operator>(const CombinedIterator& other) const {
+    return it1 > other.it1;
+  }
+  bool operator>=(const CombinedIterator& other) const {
+    return it1 >= other.it1;
+  }
+  bool operator==(const CombinedIterator& other) const {
+    return it1 == other.it1;
+  }
+  pointer operator->() const {
+    return pointer{{it1, it2}};
+  }
+  reference operator*() const {
+    return {it1, it2};
+  }
+  reference operator[](difference_type n) const {
+    return {it1 + n, it2 + n};
+  }
+
+  CombinedIterator& operator++() {
+    ++it1;
+    ++it2;
+    return *this;
+  }
+  CombinedIterator operator++(int) {
+    const auto res = *this;
+    ++*this;
+    return res;
+  }
+  CombinedIterator& operator--() {
+    --it1;
+    --it2;
+    return *this;
+  }
+  CombinedIterator operator--(int) {
+    const auto res = *this;
+    --*this;
+    return res;
+  }
+  CombinedIterator& operator+=(difference_type n) {
+    it1 += n;
+    it2 += n;
+    return *this;
+  }
+  CombinedIterator operator+(difference_type n) const {
+    CombinedIterator res = *this;
+    return res += n;
+  }
+
+  CombinedIterator& operator-=(difference_type n) {
+    it1 -= n;
+    it2 -= n;
+    return *this;
+  }
+  CombinedIterator operator-(difference_type n) const {
+    CombinedIterator res = *this;
+    return res -= n;
+  }
+  difference_type operator-(const CombinedIterator& other) {
+    return it1 - other.it1;
+  }
+
+  It1 it1;
+  It2 it2;
+};
+
+}  // namespace android
diff --git a/libs/androidfw/tests/CombinedIterator_test.cpp b/libs/androidfw/tests/CombinedIterator_test.cpp
new file mode 100644
index 0000000..c1228f3
--- /dev/null
+++ b/libs/androidfw/tests/CombinedIterator_test.cpp
@@ -0,0 +1,98 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "androidfw/CombinedIterator.h"
+
+#include <algorithm>
+#include <string>
+#include <strstream>
+#include <utility>
+#include <vector>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace android {
+
+template <class Coll>
+std::string toString(const Coll& coll) {
+  std::stringstream res;
+  res << "(" << std::size(coll) << ")";
+  if (std::size(coll)) {
+    res << "{" << coll[0];
+    for (int i = 1; i != std::size(coll); ++i) {
+      res << "," << coll[i];
+    }
+    res << "}";
+  }
+  return res.str();
+}
+
+template <class Coll>
+void AssertCollectionEq(const Coll& first, const Coll& second) {
+  ASSERT_EQ(std::size(first), std::size(second))
+      << "first: " << toString(first) << ", second: " << toString(second);
+  for (int i = 0; i != std::size(first); ++i) {
+    ASSERT_EQ(first[i], second[i])
+        << "index: " << i << " first: " << toString(first) << ", second: " << toString(second);
+  }
+}
+
+TEST(CombinedIteratorTest, Sorting) {
+  std::vector<int> v1 = {2, 1, 3, 4, 0};
+  std::vector<int> v2 = {20, 10, 30, 40, 0};
+
+  std::sort(CombinedIterator(v1.begin(), v2.begin()), CombinedIterator(v1.end(), v2.end()));
+
+  ASSERT_EQ(v1.size(), v2.size());
+  ASSERT_TRUE(std::is_sorted(v1.begin(), v1.end()));
+  ASSERT_TRUE(std::is_sorted(v2.begin(), v2.end()));
+  AssertCollectionEq(v1, {0, 1, 2, 3, 4});
+  AssertCollectionEq(v2, {0, 10, 20, 30, 40});
+}
+
+TEST(CombinedIteratorTest, Removing) {
+  std::vector<int> v1 = {1, 2, 3, 4, 5, 5, 5, 6};
+  std::vector<int> v2 = {10, 20, 30, 40, 50, 50, 50, 60};
+
+  auto newEnd =
+      std::remove_if(CombinedIterator(v1.begin(), v2.begin()), CombinedIterator(v1.end(), v2.end()),
+                     [](auto&& pair) { return pair.first >= 3 && pair.first <= 5; });
+
+  ASSERT_EQ(newEnd.it1, v1.begin() + 3);
+  ASSERT_EQ(newEnd.it2, v2.begin() + 3);
+
+  v1.erase(newEnd.it1, v1.end());
+  AssertCollectionEq(v1, {1, 2, 6});
+  v2.erase(newEnd.it2, v2.end());
+  AssertCollectionEq(v2, {10, 20, 60});
+}
+
+TEST(CombinedIteratorTest, InplaceMerge) {
+  std::vector<int> v1 = {1, 3, 4, 7, 2, 5, 6};
+  std::vector<int> v2 = {10, 30, 40, 70, 20, 50, 60};
+
+  std::inplace_merge(CombinedIterator(v1.begin(), v2.begin()),
+                     CombinedIterator(v1.begin() + 4, v2.begin() + 4),
+                     CombinedIterator(v1.end(), v2.end()));
+  ASSERT_TRUE(std::is_sorted(v1.begin(), v1.end()));
+  ASSERT_TRUE(std::is_sorted(v2.begin(), v2.end()));
+
+  AssertCollectionEq(v1, {1, 2, 3, 4, 5, 6, 7});
+  AssertCollectionEq(v2, {10, 20, 30, 40, 50, 60, 70});
+}
+
+}  // namespace android