Increase the maximum allowed dex files for boot image profiles

Boot image profiles aggregate data from a lot of apps and could easily
exceed the previous 255 dex file limit. In general a boot profile may
contain number_of_bootclasspath_elems * number_of_apps_on_device dex
files.

To cope up with this, we increase the max allowed number of dex files to
2^16. The regular profiles will keep the same low limit (255) to save ram
and space.

A side effect of this change is that in-memory representation of boot
profiles during serialization increases substantially. However, we merge
boot profiles only during idle-maintenance mode and we can afford a higher
ceiling.

Test: profile test
Bug: 139884006
Change-Id: Icdfdac8f2e4c2935692beacc8037911b64461eee
diff --git a/libprofile/profile/profile_compilation_info.cc b/libprofile/profile/profile_compilation_info.cc
index 81f1406..2e49adc 100644
--- a/libprofile/profile/profile_compilation_info.cc
+++ b/libprofile/profile/profile_compilation_info.cc
@@ -97,8 +97,8 @@
 static constexpr uint32_t kSizeWarningThresholdBytes = 500000U;
 static constexpr uint32_t kSizeErrorThresholdBytes = 1500000U;
 
-static constexpr uint32_t kSizeWarningThresholdBootBytes = 3000000U;
-static constexpr uint32_t kSizeErrorThresholdBootBytes = 6000000U;
+static constexpr uint32_t kSizeWarningThresholdBootBytes = 25000000U;
+static constexpr uint32_t kSizeErrorThresholdBootBytes = 50000000U;
 
 static bool ChecksumMatch(uint32_t dex_file_checksum, uint32_t checksum) {
   return kDebugIgnoreChecksum || dex_file_checksum == checksum;
@@ -391,8 +391,9 @@
   if (!WriteBuffer(fd, version_, sizeof(version_))) {
     return false;
   }
-  DCHECK_LE(info_.size(), std::numeric_limits<uint8_t>::max());
-  AddUintToBuffer(&buffer, static_cast<uint8_t>(info_.size()));
+
+  DCHECK_LE(info_.size(), MaxProfileIndex());
+  WriteProfileIndex(&buffer, static_cast<ProfileIndexType>(info_.size()));
 
   uint32_t required_capacity = 0;
   for (const DexFileData* dex_data_ptr : info_) {
@@ -543,7 +544,7 @@
     DCHECK_LT(classes.size(), ProfileCompilationInfo::kIndividualInlineCacheSize);
     DCHECK_NE(classes.size(), 0u) << "InlineCache contains a dex_pc with 0 classes";
 
-    SafeMap<uint8_t, std::vector<dex::TypeIndex>> dex_to_classes_map;
+    SafeMap<ProfileIndexType, std::vector<dex::TypeIndex>> dex_to_classes_map;
     // Group the classes by dex. We expect that most of the classes will come from
     // the same dex, so this will be more efficient than encoding the dex index
     // for each class reference.
@@ -551,10 +552,10 @@
     // Add the dex map size.
     AddUintToBuffer(buffer, static_cast<uint8_t>(dex_to_classes_map.size()));
     for (const auto& dex_it : dex_to_classes_map) {
-      uint8_t dex_profile_index = dex_it.first;
+      ProfileIndexType dex_profile_index = dex_it.first;
       const std::vector<dex::TypeIndex>& dex_classes = dex_it.second;
       // Add the dex profile index.
-      AddUintToBuffer(buffer, dex_profile_index);
+      WriteProfileIndex(buffer, dex_profile_index);
       // Add the the number of classes for each dex profile index.
       AddUintToBuffer(buffer, static_cast<uint8_t>(dex_classes.size()));
       for (size_t i = 0; i < dex_classes.size(); i++) {
@@ -573,11 +574,11 @@
     size += sizeof(uint16_t) * inline_cache.size();  // dex_pc
     for (const auto& inline_cache_it : inline_cache) {
       const ClassSet& classes = inline_cache_it.second.classes;
-      SafeMap<uint8_t, std::vector<dex::TypeIndex>> dex_to_classes_map;
+      SafeMap<ProfileIndexType, std::vector<dex::TypeIndex>> dex_to_classes_map;
       GroupClassesByDex(classes, &dex_to_classes_map);
       size += sizeof(uint8_t);  // dex_to_classes_map size
       for (const auto& dex_it : dex_to_classes_map) {
-        size += sizeof(uint8_t);  // dex profile index
+        size += SizeOfProfileIndexType();  // dex profile index
         size += sizeof(uint8_t);  // number of classes
         const std::vector<dex::TypeIndex>& dex_classes = dex_it.second;
         size += sizeof(uint16_t) * dex_classes.size();  // the actual classes
@@ -589,7 +590,7 @@
 
 void ProfileCompilationInfo::GroupClassesByDex(
     const ClassSet& classes,
-    /*out*/SafeMap<uint8_t, std::vector<dex::TypeIndex>>* dex_to_classes_map) {
+    /*out*/SafeMap<ProfileIndexType, std::vector<dex::TypeIndex>>* dex_to_classes_map) {
   for (const auto& classes_it : classes) {
     auto dex_it = dex_to_classes_map->FindOrAdd(classes_it.dex_profile_index);
     dex_it->second.push_back(classes_it.type_index);
@@ -601,17 +602,18 @@
     uint32_t checksum,
     uint32_t num_method_ids) {
   const auto profile_index_it = profile_key_map_.FindOrAdd(profile_key, profile_key_map_.size());
-  if (profile_key_map_.size() > std::numeric_limits<uint8_t>::max()) {
-    // Allow only 255 dex files to be profiled. This allows us to save bytes
-    // when encoding. The number is well above what we expect for normal applications.
+  if (profile_key_map_.size() > MaxProfileIndex()) {
+    // Allow only a limited number dex files to be profiled. This allows us to save bytes
+    // when encoding. For regular profiles this 2^8, and for boot profiles is 2^16
+    // (well above what we expect for normal applications).
     if (kIsDebugBuild) {
-      LOG(ERROR) << "Exceeded the maximum number of dex files (255). Something went wrong";
+      LOG(ERROR) << "Exceeded the maximum number of dex file. Something went wrong";
     }
     profile_key_map_.erase(profile_key);
     return nullptr;
   }
 
-  uint8_t profile_index = profile_index_it->second;
+  ProfileIndexType profile_index = profile_index_it->second;
   if (info_.size() <= profile_index) {
     // This is a new addition. Add it to the info_ array.
     DexFileData* dex_file_data = new (&allocator_) DexFileData(
@@ -657,7 +659,7 @@
     return nullptr;
   }
 
-  uint8_t profile_index = profile_index_it->second;
+  ProfileIndexType profile_index = profile_index_it->second;
   const DexFileData* result = info_[profile_index];
   if (verify_checksum && !ChecksumMatch(result->checksum, checksum)) {
     return nullptr;
@@ -739,8 +741,8 @@
 
 bool ProfileCompilationInfo::ReadInlineCache(
     SafeBuffer& buffer,
-    uint8_t number_of_dex_files,
-    const SafeMap<uint8_t, uint8_t>& dex_profile_index_remap,
+    ProfileIndexType number_of_dex_files,
+    const SafeMap<ProfileIndexType, ProfileIndexType>& dex_profile_index_remap,
     /*out*/ InlineCacheMap* inline_cache,
     /*out*/ std::string* error) {
   uint16_t inline_cache_size;
@@ -760,9 +762,12 @@
       continue;
     }
     for (; dex_to_classes_map_size > 0; dex_to_classes_map_size--) {
-      uint8_t dex_profile_index;
+      ProfileIndexType dex_profile_index;
       uint8_t dex_classes_size;
-      READ_UINT(uint8_t, buffer, dex_profile_index, error);
+      if (!ReadProfileIndex(buffer, &dex_profile_index)) {
+        *error = "Cannot read profile index";
+        return false;
+      }
       READ_UINT(uint8_t, buffer, dex_classes_size, error);
       if (dex_profile_index >= number_of_dex_files) {
         *error = "dex_profile_index out of bounds ";
@@ -786,11 +791,12 @@
   return true;
 }
 
-bool ProfileCompilationInfo::ReadMethods(SafeBuffer& buffer,
-                                         uint8_t number_of_dex_files,
-                                         const ProfileLineHeader& line_header,
-                                         const SafeMap<uint8_t, uint8_t>& dex_profile_index_remap,
-                                         /*out*/std::string* error) {
+bool ProfileCompilationInfo::ReadMethods(
+    SafeBuffer& buffer,
+    ProfileIndexType number_of_dex_files,
+    const ProfileLineHeader& line_header,
+    const SafeMap<ProfileIndexType, ProfileIndexType>& dex_profile_index_remap,
+    /*out*/std::string* error) {
   uint32_t unread_bytes_before_operation = buffer.CountUnreadBytes();
   if (unread_bytes_before_operation < line_header.method_region_size_bytes) {
     *error += "Profile EOF reached prematurely for ReadMethod";
@@ -919,50 +925,56 @@
 
 ProfileCompilationInfo::ProfileLoadStatus ProfileCompilationInfo::ReadProfileHeader(
       ProfileSource& source,
-      /*out*/uint8_t* number_of_dex_files,
+      /*out*/ProfileIndexType* number_of_dex_files,
       /*out*/uint32_t* uncompressed_data_size,
       /*out*/uint32_t* compressed_data_size,
       /*out*/std::string* error) {
   // Read magic and version
   const size_t kMagicVersionSize =
     sizeof(kProfileMagic) +
-    kProfileVersionSize +
-    sizeof(uint8_t) +  // number of dex files
-    sizeof(uint32_t) +  // size of uncompressed profile data
-    sizeof(uint32_t);  // size of compressed profile data
+    kProfileVersionSize;
+  SafeBuffer safe_buffer_version(kMagicVersionSize);
 
-  SafeBuffer safe_buffer(kMagicVersionSize);
-
-  ProfileLoadStatus status = safe_buffer.Fill(source, "ReadProfileHeader", error);
+  ProfileLoadStatus status = safe_buffer_version.Fill(source, "ReadProfileHeaderVersion", error);
   if (status != kProfileLoadSuccess) {
     return status;
   }
 
-  if (!safe_buffer.CompareAndAdvance(kProfileMagic, sizeof(kProfileMagic))) {
+  if (!safe_buffer_version.CompareAndAdvance(kProfileMagic, sizeof(kProfileMagic))) {
     *error = "Profile missing magic";
     return kProfileLoadVersionMismatch;
   }
-  if (safe_buffer.CountUnreadBytes() < kProfileVersionSize) {
+  if (safe_buffer_version.CountUnreadBytes() < kProfileVersionSize) {
      *error = "Cannot read profile version";
      return kProfileLoadBadData;
   }
-  memcpy(version_, safe_buffer.GetCurrentPtr(), kProfileVersionSize);
-  safe_buffer.Advance(kProfileVersionSize);
+  memcpy(version_, safe_buffer_version.GetCurrentPtr(), kProfileVersionSize);
   if ((memcmp(version_, kProfileVersion, kProfileVersionSize) != 0) &&
       (memcmp(version_, kProfileVersionForBootImage, kProfileVersionSize) != 0)) {
     *error = "Profile version mismatch";
     return kProfileLoadVersionMismatch;
   }
 
-  if (!safe_buffer.ReadUintAndAdvance<uint8_t>(number_of_dex_files)) {
+  const size_t kProfileHeaderDataSize =
+    SizeOfProfileIndexType() +  // number of dex files
+    sizeof(uint32_t) +  // size of uncompressed profile data
+    sizeof(uint32_t);  // size of compressed profile data
+  SafeBuffer safe_buffer_header_data(kProfileHeaderDataSize);
+
+  status = safe_buffer_header_data.Fill(source, "ReadProfileHeaderData", error);
+  if (status != kProfileLoadSuccess) {
+    return status;
+  }
+
+  if (!ReadProfileIndex(safe_buffer_header_data, number_of_dex_files)) {
     *error = "Cannot read the number of dex files";
     return kProfileLoadBadData;
   }
-  if (!safe_buffer.ReadUintAndAdvance<uint32_t>(uncompressed_data_size)) {
+  if (!safe_buffer_header_data.ReadUintAndAdvance<uint32_t>(uncompressed_data_size)) {
     *error = "Cannot read the size of uncompressed data";
     return kProfileLoadBadData;
   }
-  if (!safe_buffer.ReadUintAndAdvance<uint32_t>(compressed_data_size)) {
+  if (!safe_buffer_header_data.ReadUintAndAdvance<uint32_t>(compressed_data_size)) {
     *error = "Cannot read the size of compressed data";
     return kProfileLoadBadData;
   }
@@ -1014,9 +1026,9 @@
 
 ProfileCompilationInfo::ProfileLoadStatus ProfileCompilationInfo::ReadProfileLine(
       SafeBuffer& buffer,
-      uint8_t number_of_dex_files,
+      ProfileIndexType number_of_dex_files,
       const ProfileLineHeader& line_header,
-      const SafeMap<uint8_t, uint8_t>& dex_profile_index_remap,
+      const SafeMap<ProfileIndexType, ProfileIndexType>& dex_profile_index_remap,
       bool merge_classes,
       /*out*/std::string* error) {
   DexFileData* data = GetOrAddDexFileData(line_header.profile_key,
@@ -1117,13 +1129,13 @@
         }
 
         const ClassSet &classes = dex_pc_data.classes;
-        SafeMap<uint8_t, std::vector<dex::TypeIndex>> dex_to_classes_map;
+        SafeMap<ProfileIndexType, std::vector<dex::TypeIndex>> dex_to_classes_map;
         // Group the classes by dex. We expect that most of the classes will come from
         // the same dex, so this will be more efficient than encoding the dex index
         // for each class reference.
         GroupClassesByDex(classes, &dex_to_classes_map);
         for (const auto &dex_it : dex_to_classes_map) {
-          uint8_t dex_profile_index = dex_it.first;
+          ProfileIndexType dex_profile_index = dex_it.first;
           const auto dex_file_inline_cache_it = key_to_dex_file.find(
               info_[dex_profile_index]->profile_key);
           if (dex_file_inline_cache_it == key_to_dex_file.end()) {
@@ -1273,7 +1285,7 @@
   }
 
   // Read profile header: magic + version + number_of_dex_files.
-  uint8_t number_of_dex_files;
+  ProfileIndexType number_of_dex_files;
   uint32_t uncompressed_data_size;
   uint32_t compressed_data_size;
   status = ReadProfileHeader(*source,
@@ -1325,7 +1337,7 @@
 
   std::vector<ProfileLineHeader> profile_line_headers;
   // Read profile line headers.
-  for (uint8_t k = 0; k < number_of_dex_files; k++) {
+  for (ProfileIndexType k = 0; k < number_of_dex_files; k++) {
     ProfileLineHeader line_header;
 
     // First, read the line header to get the amount of data we need to read.
@@ -1336,12 +1348,12 @@
     profile_line_headers.push_back(line_header);
   }
 
-  SafeMap<uint8_t, uint8_t> dex_profile_index_remap;
+  SafeMap<ProfileIndexType, ProfileIndexType> dex_profile_index_remap;
   if (!RemapProfileIndex(profile_line_headers, filter_fn, &dex_profile_index_remap)) {
     return kProfileLoadBadData;
   }
 
-  for (uint8_t k = 0; k < number_of_dex_files; k++) {
+  for (ProfileIndexType k = 0; k < number_of_dex_files; k++) {
     if (!filter_fn(profile_line_headers[k].profile_key, profile_line_headers[k].checksum)) {
       // We have to skip the line. Advanced the current pointer of the buffer.
       size_t profile_line_size =
@@ -1377,7 +1389,7 @@
 bool ProfileCompilationInfo::RemapProfileIndex(
     const std::vector<ProfileLineHeader>& profile_line_headers,
     const ProfileLoadFilterFn& filter_fn,
-    /*out*/SafeMap<uint8_t, uint8_t>* dex_profile_index_remap) {
+    /*out*/SafeMap<ProfileIndexType, ProfileIndexType>* dex_profile_index_remap) {
   // First verify that all checksums match. This will avoid adding garbage to
   // the current profile info.
   // Note that the number of elements should be very small, so this should not
@@ -1501,7 +1513,7 @@
 
   // First, build a mapping from other_dex_profile_index to this_dex_profile_index.
   // This will make sure that the ClassReferences  will point to the correct dex file.
-  SafeMap<uint8_t, uint8_t> dex_profile_index_remap;
+  SafeMap<ProfileIndexType, ProfileIndexType> dex_profile_index_remap;
   for (const DexFileData* other_dex_data : other.info_) {
     const DexFileData* dex_data = GetOrAddDexFileData(other_dex_data->profile_key,
                                                       other_dex_data->checksum,
@@ -2215,4 +2227,37 @@
       const ProfileSampleAnnotation& other) const {
   return origin_package_name_ == other.origin_package_name_;
 }
+
+void ProfileCompilationInfo::WriteProfileIndex(
+    std::vector<uint8_t>* buffer, ProfileIndexType value) const {
+  if (IsForBootImage()) {
+    AddUintToBuffer(buffer, value);
+  } else {
+    AddUintToBuffer(buffer, static_cast<ProfileIndexTypeRegular>(value));
+  }
+}
+
+bool ProfileCompilationInfo::ReadProfileIndex(
+    SafeBuffer& safe_buffer, ProfileIndexType* value) const {
+  if (IsForBootImage()) {
+    return safe_buffer.ReadUintAndAdvance<ProfileIndexType>(value);
+  } else {
+    ProfileIndexTypeRegular out;
+    bool result = safe_buffer.ReadUintAndAdvance<ProfileIndexTypeRegular>(&out);
+    *value = out;
+    return result;
+  }
+}
+
+ProfileCompilationInfo::ProfileIndexType ProfileCompilationInfo::MaxProfileIndex() const {
+  return IsForBootImage()
+      ? std::numeric_limits<ProfileIndexType>::max()
+      : std::numeric_limits<ProfileIndexTypeRegular>::max();
+}
+
+uint32_t ProfileCompilationInfo::SizeOfProfileIndexType() const {
+  return IsForBootImage()
+    ? sizeof(ProfileIndexType)
+    : sizeof(ProfileIndexTypeRegular);
+}
 }  // namespace art
diff --git a/libprofile/profile/profile_compilation_info.h b/libprofile/profile/profile_compilation_info.h
index 52f42f5..90bf069 100644
--- a/libprofile/profile/profile_compilation_info.h
+++ b/libprofile/profile/profile_compilation_info.h
@@ -106,6 +106,21 @@
     uint32_t num_method_ids;
   };
 
+  // The types used to manipulate the profile index of dex files.
+  // They set an upper limit to how many dex files a given profile can recored.
+  //
+  // Boot profiles have more needs than regular profiles as they contain data from
+  // many apps merged together. As such they set the default type for data manipulation.
+  //
+  // Regular profiles don't record a lot of dex files, and use a smaller data type
+  // in order to save disk and ram.
+  //
+  // In-memory all profiles will use ProfileIndexType to represent the indices. However,
+  // when serialized, the profile type (boot or regular) will determine which data type
+  // is used to write the data.
+  using ProfileIndexType = uint16_t;
+  using ProfileIndexTypeRegular = uint8_t;
+
   // Encodes a class reference in the profile.
   // The owning dex file is encoded as the index (dex_profile_index) it has in the
   // profile rather than as a full DexRefence(location,checksum).
@@ -118,7 +133,7 @@
   // data from multiple splits. This means that a profile may contain a classes2.dex from split-A
   // and one from split-B.
   struct ClassReference : public ValueObject {
-    ClassReference(uint8_t dex_profile_idx, const dex::TypeIndex type_idx) :
+    ClassReference(ProfileIndexType dex_profile_idx, const dex::TypeIndex type_idx) :
       dex_profile_index(dex_profile_idx), type_index(type_idx) {}
 
     bool operator==(const ClassReference& other) const {
@@ -130,7 +145,7 @@
           : dex_profile_index < other.dex_profile_index;
     }
 
-    uint8_t dex_profile_index;  // the index of the owning dex in the profile info
+    ProfileIndexType dex_profile_index;  // the index of the owning dex in the profile info
     dex::TypeIndex type_index;  // the type index of the class
   };
 
@@ -629,7 +644,7 @@
     // The profile key this data belongs to.
     std::string profile_key;
     // The profile index of this dex file (matches ClassReference#dex_profile_index).
-    uint8_t profile_index;
+    ProfileIndexType profile_index;
     // The dex checksum.
     uint32_t checksum;
     // The methods' profile information.
@@ -813,7 +828,7 @@
   // Read the profile header from the given fd and store the number of profile
   // lines into number_of_dex_files.
   ProfileLoadStatus ReadProfileHeader(ProfileSource& source,
-                                      /*out*/uint8_t* number_of_dex_files,
+                                      /*out*/ProfileIndexType* number_of_dex_files,
                                       /*out*/uint32_t* size_uncompressed_data,
                                       /*out*/uint32_t* size_compressed_data,
                                       /*out*/std::string* error);
@@ -830,12 +845,13 @@
                                      /*out*/std::string* error);
 
   // Read a single profile line from the given fd.
-  ProfileLoadStatus ReadProfileLine(SafeBuffer& buffer,
-                                    uint8_t number_of_dex_files,
-                                    const ProfileLineHeader& line_header,
-                                    const SafeMap<uint8_t, uint8_t>& dex_profile_index_remap,
-                                    bool merge_classes,
-                                    /*out*/std::string* error);
+  ProfileLoadStatus ReadProfileLine(
+      SafeBuffer& buffer,
+      ProfileIndexType number_of_dex_files,
+      const ProfileLineHeader& line_header,
+      const SafeMap<ProfileIndexType, ProfileIndexType>& dex_profile_index_remap,
+      bool merge_classes,
+      /*out*/std::string* error);
 
   // Read all the classes from the buffer into the profile `info_` structure.
   bool ReadClasses(SafeBuffer& buffer,
@@ -844,21 +860,22 @@
 
   // Read all the methods from the buffer into the profile `info_` structure.
   bool ReadMethods(SafeBuffer& buffer,
-                   uint8_t number_of_dex_files,
+                   ProfileIndexType number_of_dex_files,
                    const ProfileLineHeader& line_header,
-                   const SafeMap<uint8_t, uint8_t>& dex_profile_index_remap,
+                   const SafeMap<ProfileIndexType, ProfileIndexType>& dex_profile_index_remap,
                    /*out*/std::string* error);
 
   // The method generates mapping of profile indices while merging a new profile
   // data into current data. It returns true, if the mapping was successful.
-  bool RemapProfileIndex(const std::vector<ProfileLineHeader>& profile_line_headers,
-                         const ProfileLoadFilterFn& filter_fn,
-                         /*out*/SafeMap<uint8_t, uint8_t>* dex_profile_index_remap);
+  bool RemapProfileIndex(
+      const std::vector<ProfileLineHeader>& profile_line_headers,
+      const ProfileLoadFilterFn& filter_fn,
+      /*out*/SafeMap<ProfileIndexType, ProfileIndexType>* dex_profile_index_remap);
 
   // Read the inline cache encoding from line_bufer into inline_cache.
   bool ReadInlineCache(SafeBuffer& buffer,
-                       uint8_t number_of_dex_files,
-                       const SafeMap<uint8_t, uint8_t>& dex_profile_index_remap,
+                       ProfileIndexType number_of_dex_files,
+                       const SafeMap<ProfileIndexType, ProfileIndexType>& dex_profile_index_remap,
                        /*out*/InlineCacheMap* inline_cache,
                        /*out*/std::string* error);
 
@@ -874,7 +891,7 @@
   // `dex_to_classes_map`.
   void GroupClassesByDex(
       const ClassSet& classes,
-      /*out*/SafeMap<uint8_t, std::vector<dex::TypeIndex>>* dex_to_classes_map);
+      /*out*/SafeMap<ProfileIndexType, std::vector<dex::TypeIndex>>* dex_to_classes_map);
 
   // Find the data for the dex_pc in the inline cache. Adds an empty entry
   // if no previous data exists.
@@ -900,6 +917,18 @@
   static std::string MigrateAnnotationInfo(const std::string& base_key,
                                            const std::string& augmented_key);
 
+  // Returns the maximum value for the profile index. It depends on the profile type.
+  // Boot profiles can store more dex files than regular profiles.
+  ProfileIndexType MaxProfileIndex() const;
+  // Returns the size of the profile index type used for serialization.
+  uint32_t SizeOfProfileIndexType() const;
+  // Writes the profile index to the buffer. The type of profile will determine the
+  // number of bytes used for serialization.
+  void WriteProfileIndex(std::vector<uint8_t>* buffer, ProfileIndexType value) const;
+  // Read the profile index from the buffer. The type of profile will determine the
+  // number of bytes used for serialization.
+  bool ReadProfileIndex(SafeBuffer& safe_buffer, ProfileIndexType* value) const;
+
   friend class ProfileCompilationInfoTest;
   friend class CompilerDriverProfileTest;
   friend class ProfileAssistantTest;
@@ -916,7 +945,7 @@
   // Cache mapping profile keys to profile index.
   // This is used to speed up searches since it avoids iterating
   // over the info_ vector when searching by profile key.
-  ArenaSafeMap<const std::string, uint8_t> profile_key_map_;
+  ArenaSafeMap<const std::string, ProfileIndexType> profile_key_map_;
 
   // The version of the profile.
   uint8_t version_[kProfileVersionSize];
diff --git a/libprofile/profile/profile_compilation_info_test.cc b/libprofile/profile/profile_compilation_info_test.cc
index a6ef146..8b1cedf 100644
--- a/libprofile/profile/profile_compilation_info_test.cc
+++ b/libprofile/profile/profile_compilation_info_test.cc
@@ -34,6 +34,8 @@
 using Hotness = ProfileCompilationInfo::MethodHotness;
 using ProfileInlineCache = ProfileMethodInfo::ProfileInlineCache;
 using ProfileSampleAnnotation = ProfileCompilationInfo::ProfileSampleAnnotation;
+using ProfileIndexType = ProfileCompilationInfo::ProfileIndexType;
+using ProfileIndexTypeRegular = ProfileCompilationInfo::ProfileIndexTypeRegular;
 
 static constexpr size_t kMaxMethodIds = 65535;
 static uint32_t kMaxHotnessFlagBootIndex =
@@ -689,18 +691,33 @@
   }
 }
 
-TEST_F(ProfileCompilationInfoTest, AddMoreDexFileThanLimit) {
+TEST_F(ProfileCompilationInfoTest, AddMoreDexFileThanLimitRegular) {
   FakeDexStorage local_storage;
   ProfileCompilationInfo info;
   // Save a few methods.
-  for (uint16_t i = 0; i < std::numeric_limits<uint8_t>::max(); i++) {
+  for (uint16_t i = 0; i < std::numeric_limits<ProfileIndexTypeRegular>::max(); i++) {
     std::string location = std::to_string(i);
     const DexFile* dex = local_storage.AddFakeDex(
         location, /* checksum= */ 1, /* num_method_ids= */ 1);
     ASSERT_TRUE(AddMethod(&info, dex, /* method_idx= */ 0));
   }
-  // We only support at most 255 dex files.
-  const DexFile* dex = local_storage.AddFakeDex("256", /* checksum= */ 1, /* num_method_ids= */ 1);
+  // Add an extra dex file.
+  const DexFile* dex = local_storage.AddFakeDex("-1", /* checksum= */ 1, /* num_method_ids= */ 1);
+  ASSERT_FALSE(AddMethod(&info, dex, /* method_idx= */ 0));
+}
+
+TEST_F(ProfileCompilationInfoTest, AddMoreDexFileThanLimitBoot) {
+  FakeDexStorage local_storage;
+  ProfileCompilationInfo info(/*for_boot_image=*/true);
+  // Save a few methods.
+  for (uint16_t i = 0; i < std::numeric_limits<ProfileIndexType>::max(); i++) {
+    std::string location = std::to_string(i);
+    const DexFile* dex = local_storage.AddFakeDex(
+        location, /* checksum= */ 1, /* num_method_ids= */ 1);
+    ASSERT_TRUE(AddMethod(&info, dex, /* method_idx= */ 0));
+  }
+  // Add an extra dex file.
+  const DexFile* dex = local_storage.AddFakeDex("-1", /* checksum= */ 1, /* num_method_ids= */ 1);
   ASSERT_FALSE(AddMethod(&info, dex, /* method_idx= */ 0));
 }