Merge "[res] Optimize idmap format for lookups" into main

6 files changed, 124 insertions, 89 deletions

diff --git a/cmds/idmap2/include/idmap2/Idmap.h b/cmds/idmap2/include/idmap2/Idmap.h
index e86f81485a41..b0ba01957ab6 100644
--- a/cmds/idmap2/include/idmap2/Idmap.h
+++ b/cmds/idmap2/include/idmap2/Idmap.h
@@ -21,18 +21,19 @@
  * header                     := magic version target_crc overlay_crc fulfilled_policies
  *                               enforce_overlayable target_path overlay_path overlay_name
  *                               debug_info
- * data                       := data_header target_entry* target_inline_entry*
-                                 target_inline_entry_value* config* overlay_entry* string_pool
+ * data                       := data_header target_entries target_inline_entries
+                                 target_inline_entry_value* config* overlay_entries string_pool
  * data_header                := target_entry_count target_inline_entry_count
                                  target_inline_entry_value_count config_count overlay_entry_count
  *                               string_pool_index
- * target_entry               := target_id overlay_id
- * target_inline_entry        := target_id start_value_index value_count
+ * target_entries             := target_id* overlay_id*
+ * target_inline_entries      := target_id* target_inline_value_header*
+ * target_inline_value_header := start_value_index value_count
  * target_inline_entry_value  := config_index Res_value::size padding(1) Res_value::type
  *                               Res_value::value
  * config                     := target_id Res_value::size padding(1) Res_value::type
  *                               Res_value::value
- * overlay_entry              := overlay_id target_id
+ * overlay_entries            := overlay_id* target_id*
  *
  * debug_info                       := string
  * enforce_overlayable              := <uint32_t>
diff --git a/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp b/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp
index 89769246434a..00ef0c7f8cf0 100644
--- a/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp
+++ b/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp
@@ -66,43 +66,57 @@ void BinaryStreamVisitor::visit(const IdmapHeader& header) {
 void BinaryStreamVisitor::visit(const IdmapData& data) {
   for (const auto& target_entry : data.GetTargetEntries()) {
     Write32(target_entry.target_id);
+  }
+  for (const auto& target_entry : data.GetTargetEntries()) {
     Write32(target_entry.overlay_id);
   }
 
-  static constexpr uint16_t kValueSize = 8U;
-  std::vector<std::pair<ConfigDescription, TargetValue>> target_values;
-  target_values.reserve(data.GetHeader()->GetTargetInlineEntryValueCount());
-  for (const auto& target_entry : data.GetTargetInlineEntries()) {
-    Write32(target_entry.target_id);
-    Write32(target_values.size());
-    Write32(target_entry.values.size());
-    target_values.insert(
-        target_values.end(), target_entry.values.begin(), target_entry.values.end());
+  uint32_t current_inline_entry_values_count = 0;
+  for (const auto& target_inline_entry : data.GetTargetInlineEntries()) {
+    Write32(target_inline_entry.target_id);
+  }
+  for (const auto& target_inline_entry : data.GetTargetInlineEntries()) {
+    Write32(current_inline_entry_values_count);
+    Write32(target_inline_entry.values.size());
+    current_inline_entry_values_count += target_inline_entry.values.size();
   }
 
   std::vector<ConfigDescription> configs;
   configs.reserve(data.GetHeader()->GetConfigCount());
-  for (const auto& target_entry_value : target_values) {
-    auto config_it = find(configs.begin(), configs.end(), target_entry_value.first);
-    if (config_it != configs.end()) {
-      Write32(config_it - configs.begin());
-    } else {
-      Write32(configs.size());
-      configs.push_back(target_entry_value.first);
+  for (const auto& target_entry : data.GetTargetInlineEntries()) {
+    for (const auto& target_entry_value : target_entry.values) {
+      auto config_it = std::find(configs.begin(), configs.end(), target_entry_value.first);
+      if (config_it != configs.end()) {
+        Write32(config_it - configs.begin());
+      } else {
+        Write32(configs.size());
+        configs.push_back(target_entry_value.first);
+      }
+      // We're writing a Res_value entry here, and the first 3 bytes of that are
+      // sizeof() + a padding 0 byte
+      static constexpr decltype(android::Res_value::size) kSize = sizeof(android::Res_value);
+      Write16(kSize);
+      Write8(0U);
+      Write8(target_entry_value.second.data_type);
+      Write32(target_entry_value.second.data_value);
     }
-    Write16(kValueSize);
-    Write8(0U);  // padding
-    Write8(target_entry_value.second.data_type);
-    Write32(target_entry_value.second.data_value);
   }
 
-  for( auto& cd : configs) {
-    cd.swapHtoD();
-    stream_.write(reinterpret_cast<char*>(&cd), sizeof(cd));
+  if (!configs.empty()) {
+    stream_.write(reinterpret_cast<const char*>(&configs.front()),
+                  sizeof(configs.front()) * configs.size());
+    if (configs.size() >= 100) {
+      // Let's write a message to future us so that they know when to replace the linear search
+      // in `configs` vector with something more efficient.
+      LOG(WARNING) << "Idmap got " << configs.size()
+                   << " configurations, time to fix the bruteforce search";
+    }
   }
 
   for (const auto& overlay_entry : data.GetOverlayEntries()) {
     Write32(overlay_entry.overlay_id);
+  }
+  for (const auto& overlay_entry : data.GetOverlayEntries()) {
     Write32(overlay_entry.target_id);
   }
 
diff --git a/cmds/idmap2/libidmap2/Idmap.cpp b/cmds/idmap2/libidmap2/Idmap.cpp
index 12d9dd9bd1ad..7680109f1d54 100644
--- a/cmds/idmap2/libidmap2/Idmap.cpp
+++ b/cmds/idmap2/libidmap2/Idmap.cpp
@@ -204,73 +204,91 @@ std::unique_ptr<const IdmapData> IdmapData::FromBinaryStream(std::istream& strea
   }
 
   // Read the mapping of target resource id to overlay resource value.
+  data->target_entries_.resize(data->header_->GetTargetEntryCount());
   for (size_t i = 0; i < data->header_->GetTargetEntryCount(); i++) {
-    TargetEntry target_entry{};
-    if (!Read32(stream, &target_entry.target_id) || !Read32(stream, &target_entry.overlay_id)) {
+    if (!Read32(stream, &data->target_entries_[i].target_id)) {
+      return nullptr;
+    }
+  }
+  for (size_t i = 0; i < data->header_->GetTargetEntryCount(); i++) {
+    if (!Read32(stream, &data->target_entries_[i].overlay_id)) {
       return nullptr;
     }
-    data->target_entries_.emplace_back(target_entry);
   }
 
   // Read the mapping of target resource id to inline overlay values.
-  std::vector<std::tuple<TargetInlineEntry, uint32_t, uint32_t>> target_inline_entries;
+  struct TargetInlineEntryHeader {
+    ResourceId target_id;
+    uint32_t values_offset;
+    uint32_t values_count;
+  };
+  std::vector<TargetInlineEntryHeader> target_inline_entries(
+      data->header_->GetTargetInlineEntryCount());
+  for (size_t i = 0; i < data->header_->GetTargetInlineEntryCount(); i++) {
+    if (!Read32(stream, &target_inline_entries[i].target_id)) {
+      return nullptr;
+    }
+  }
   for (size_t i = 0; i < data->header_->GetTargetInlineEntryCount(); i++) {
-    TargetInlineEntry target_entry{};
-    uint32_t entry_offset;
-    uint32_t entry_count;
-    if (!Read32(stream, &target_entry.target_id) || !Read32(stream, &entry_offset)
-        || !Read32(stream, &entry_count)) {
+    if (!Read32(stream, &target_inline_entries[i].values_offset) ||
+        !Read32(stream, &target_inline_entries[i].values_count)) {
       return nullptr;
     }
-    target_inline_entries.emplace_back(target_entry, entry_offset, entry_count);
   }
 
   // Read the inline overlay resource values
-  std::vector<std::pair<uint32_t, TargetValue>> target_values;
-  uint8_t unused1;
-  uint16_t unused2;
-  for (size_t i = 0; i < data->header_->GetTargetInlineEntryValueCount(); i++) {
+  struct TargetValueHeader {
     uint32_t config_index;
-    if (!Read32(stream, &config_index)) {
+    DataType data_type;
+    DataValue data_value;
+  };
+  std::vector<TargetValueHeader> target_values(data->header_->GetTargetInlineEntryValueCount());
+  for (size_t i = 0; i < data->header_->GetTargetInlineEntryValueCount(); i++) {
+    auto& value = target_values[i];
+    if (!Read32(stream, &value.config_index)) {
       return nullptr;
     }
-    TargetValue value;
-    if (!Read16(stream, &unused2)
-        || !Read8(stream, &unused1)
-        || !Read8(stream, &value.data_type)
-        || !Read32(stream, &value.data_value)) {
+    // skip the padding
+    stream.seekg(3, std::ios::cur);
+    if (!Read8(stream, &value.data_type) || !Read32(stream, &value.data_value)) {
       return nullptr;
     }
-    target_values.emplace_back(config_index, value);
   }
 
   // Read the configurations
-  std::vector<ConfigDescription> configurations;
-  for (size_t i = 0; i < data->header_->GetConfigCount(); i++) {
-    ConfigDescription cd;
-    if (!stream.read(reinterpret_cast<char*>(&cd), sizeof(ConfigDescription))) {
+  std::vector<ConfigDescription> configurations(data->header_->GetConfigCount());
+  if (!configurations.empty()) {
+    if (!stream.read(reinterpret_cast<char*>(&configurations.front()),
+                     sizeof(configurations.front()) * configurations.size())) {
       return nullptr;
     }
-    configurations.emplace_back(cd);
   }
 
   // Construct complete target inline entries
-  for (auto [target_entry, entry_offset, entry_count] : target_inline_entries) {
-    for(size_t i = 0; i < entry_count; i++) {
-      const auto& target_value = target_values[entry_offset + i];
-      const auto& config = configurations[target_value.first];
-      target_entry.values[config] = target_value.second;
+  data->target_inline_entries_.reserve(target_inline_entries.size());
+  for (auto&& entry_header : target_inline_entries) {
+    TargetInlineEntry& entry = data->target_inline_entries_.emplace_back();
+    entry.target_id = entry_header.target_id;
+    for (size_t i = 0; i < entry_header.values_count; i++) {
+      const auto& value_header = target_values[entry_header.values_offset + i];
+      const auto& config = configurations[value_header.config_index];
+      auto& value = entry.values[config];
+      value.data_type = value_header.data_type;
+      value.data_value = value_header.data_value;
     }
-    data->target_inline_entries_.emplace_back(target_entry);
   }
 
   // Read the mapping of overlay resource id to target resource id.
+  data->overlay_entries_.resize(data->header_->GetOverlayEntryCount());
+  for (size_t i = 0; i < data->header_->GetOverlayEntryCount(); i++) {
+    if (!Read32(stream, &data->overlay_entries_[i].overlay_id)) {
+      return nullptr;
+    }
+  }
   for (size_t i = 0; i < data->header_->GetOverlayEntryCount(); i++) {
-    OverlayEntry overlay_entry{};
-    if (!Read32(stream, &overlay_entry.overlay_id) || !Read32(stream, &overlay_entry.target_id)) {
+    if (!Read32(stream, &data->overlay_entries_[i].target_id)) {
       return nullptr;
     }
-    data->overlay_entries_.emplace_back(overlay_entry);
   }
 
   // Read raw string pool bytes.
@@ -320,7 +338,7 @@ Result<std::unique_ptr<const IdmapData>> IdmapData::FromResourceMapping(
   std::unique_ptr<IdmapData> data(new IdmapData());
   data->string_pool_data_ = std::string(resource_mapping.GetStringPoolData());
   uint32_t inline_value_count = 0;
-  std::set<std::string> config_set;
+  std::set<std::string_view> config_set;
   for (const auto& mapping : resource_mapping.GetTargetToOverlayMap()) {
     if (auto overlay_resource = std::get_if<ResourceId>(&mapping.second)) {
       data->target_entries_.push_back({mapping.first, *overlay_resource});
@@ -329,7 +347,9 @@ Result<std::unique_ptr<const IdmapData>> IdmapData::FromResourceMapping(
       for (const auto& [config, value] : std::get<ConfigMap>(mapping.second)) {
         config_set.insert(config);
         ConfigDescription cd;
-        ConfigDescription::Parse(config, &cd);
+        if (!ConfigDescription::Parse(config, &cd)) {
+          return Error("failed to parse configuration string '%s'", config.c_str());
+        }
         values[cd] = value;
         inline_value_count++;
       }
diff --git a/cmds/idmap2/tests/IdmapTests.cpp b/cmds/idmap2/tests/IdmapTests.cpp
index c85619c1e4bf..1b656e8c2088 100644
--- a/cmds/idmap2/tests/IdmapTests.cpp
+++ b/cmds/idmap2/tests/IdmapTests.cpp
@@ -68,7 +68,7 @@ TEST(IdmapTests, CreateIdmapHeaderFromBinaryStream) {
   std::unique_ptr<const IdmapHeader> header = IdmapHeader::FromBinaryStream(stream);
   ASSERT_THAT(header, NotNull());
   ASSERT_EQ(header->GetMagic(), 0x504d4449U);
-  ASSERT_EQ(header->GetVersion(), 0x09U);
+  ASSERT_EQ(header->GetVersion(), 10);
   ASSERT_EQ(header->GetTargetCrc(), 0x1234U);
   ASSERT_EQ(header->GetOverlayCrc(), 0x5678U);
   ASSERT_EQ(header->GetFulfilledPolicies(), 0x11);
@@ -143,7 +143,7 @@ TEST(IdmapTests, CreateIdmapFromBinaryStream) {
 
   ASSERT_THAT(idmap->GetHeader(), NotNull());
   ASSERT_EQ(idmap->GetHeader()->GetMagic(), 0x504d4449U);
-  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 0x09U);
+  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 10);
   ASSERT_EQ(idmap->GetHeader()->GetTargetCrc(), 0x1234U);
   ASSERT_EQ(idmap->GetHeader()->GetOverlayCrc(), 0x5678U);
   ASSERT_EQ(idmap->GetHeader()->GetFulfilledPolicies(), kIdmapRawDataPolicies);
@@ -204,7 +204,7 @@ TEST(IdmapTests, CreateIdmapHeaderFromApkAssets) {
 
   ASSERT_THAT(idmap->GetHeader(), NotNull());
   ASSERT_EQ(idmap->GetHeader()->GetMagic(), 0x504d4449U);
-  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 0x09U);
+  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 10);
   ASSERT_EQ(idmap->GetHeader()->GetTargetCrc(), android::idmap2::TestConstants::TARGET_CRC);
   ASSERT_EQ(idmap->GetHeader()->GetOverlayCrc(), android::idmap2::TestConstants::OVERLAY_CRC);
   ASSERT_EQ(idmap->GetHeader()->GetFulfilledPolicies(), PolicyFlags::PUBLIC);
diff --git a/cmds/idmap2/tests/RawPrintVisitorTests.cpp b/cmds/idmap2/tests/RawPrintVisitorTests.cpp
index 68164e26f352..7fae1c64f014 100644
--- a/cmds/idmap2/tests/RawPrintVisitorTests.cpp
+++ b/cmds/idmap2/tests/RawPrintVisitorTests.cpp
@@ -64,7 +64,7 @@ TEST(RawPrintVisitorTests, CreateRawPrintVisitor) {
   (*idmap)->accept(&visitor);
 
   ASSERT_CONTAINS_REGEX(ADDRESS "504d4449  magic\n", stream.str());
-  ASSERT_CONTAINS_REGEX(ADDRESS "00000009  version\n", stream.str());
+  ASSERT_CONTAINS_REGEX(ADDRESS "0000000a  version\n", stream.str());
   ASSERT_CONTAINS_REGEX(
       StringPrintf(ADDRESS "%s  target crc\n", android::idmap2::TestConstants::TARGET_CRC_STRING),
       stream.str());
@@ -113,7 +113,7 @@ TEST(RawPrintVisitorTests, CreateRawPrintVisitorWithoutAccessToApks) {
   (*idmap)->accept(&visitor);
 
   ASSERT_CONTAINS_REGEX(ADDRESS "504d4449  magic\n", stream.str());
-  ASSERT_CONTAINS_REGEX(ADDRESS "00000009  version\n", stream.str());
+  ASSERT_CONTAINS_REGEX(ADDRESS "0000000a  version\n", stream.str());
   ASSERT_CONTAINS_REGEX(ADDRESS "00001234  target crc\n", stream.str());
   ASSERT_CONTAINS_REGEX(ADDRESS "00005678  overlay crc\n", stream.str());
   ASSERT_CONTAINS_REGEX(ADDRESS "00000011  fulfilled policies: public|signature\n", stream.str());
diff --git a/cmds/idmap2/tests/TestHelpers.h b/cmds/idmap2/tests/TestHelpers.h
index bf01c32c4c86..2b4ebd1ae800 100644
--- a/cmds/idmap2/tests/TestHelpers.h
+++ b/cmds/idmap2/tests/TestHelpers.h
@@ -34,7 +34,7 @@ const unsigned char kIdmapRawData[] = {
     0x49, 0x44, 0x4d, 0x50,
 
     // 0x4: version
-    0x09, 0x00, 0x00, 0x00,
+    0x0a, 0x00, 0x00, 0x00,
 
     // 0x8: target crc
     0x34, 0x12, 0x00, 0x00,
@@ -95,19 +95,15 @@ const unsigned char kIdmapRawData[] = {
     // TARGET ENTRIES
     // 0x6c: target id (0x7f020000)
     0x00, 0x00, 0x02, 0x7f,
-
-    // 0x70: overlay_id (0x7f020000)
-    0x00, 0x00, 0x02, 0x7f,
-
-    // 0x74: target id (0x7f030000)
-    0x00, 0x00, 0x03, 0x7f,
-
-    // 0x78: overlay_id (0x7f030000)
+    // 0x70: target id (0x7f030000)
     0x00, 0x00, 0x03, 0x7f,
-
-    // 0x7c: target id (0x7f030002)
+    // 0x74: target id (0x7f030002)
     0x02, 0x00, 0x03, 0x7f,
 
+    // 0x78: overlay_id (0x7f020000)
+    0x00, 0x00, 0x02, 0x7f,
+    // 0x7c: overlay_id (0x7f030000)
+    0x00, 0x00, 0x03, 0x7f,
     // 0x80: overlay_id (0x7f030001)
     0x01, 0x00, 0x03, 0x7f,
 
@@ -178,16 +174,20 @@ const unsigned char kIdmapRawData[] = {
     // 0xe1: padding
     0x00, 0x00, 0x00,
 
-
     // OVERLAY ENTRIES
-    // 0xe4: 0x7f020000 -> 0x7f020000
-    0x00, 0x00, 0x02, 0x7f, 0x00, 0x00, 0x02, 0x7f,
-
-    // 0xec: 0x7f030000 -> 0x7f030000
-    0x00, 0x00, 0x03, 0x7f, 0x00, 0x00, 0x03, 0x7f,
+    // 0xe4: 0x7f020000 -> ...
+    0x00, 0x00, 0x02, 0x7f,
+    // 0xe8: 0x7f030000 -> ...
+    0x00, 0x00, 0x03, 0x7f,
+    // 0xec: 0x7f030001 -> ...
+    0x01, 0x00, 0x03, 0x7f,
 
-    // 0xf4: 0x7f030001 -> 0x7f030002
-    0x01, 0x00, 0x03, 0x7f, 0x02, 0x00, 0x03, 0x7f,
+    // 0xf0: ... -> 0x7f020000
+    0x00, 0x00, 0x02, 0x7f,
+    // 0xf4: ... -> 0x7f030000
+    0x00, 0x00, 0x03, 0x7f,
+    // 0xf8: ... -> 0x7f030002
+    0x02, 0x00, 0x03, 0x7f,
 
     // 0xfc: string pool
     // string length,