Fix ProfileCompilationInfo::UpdateProfileKeys().

Correctly migrate the annotation.

Also do some cleanup. Use `std::unique_ptr<>` for `info_`.
Change the `profile_key_map_` key to `std::string_view`,
using the `DexFileData::profile_key` as backing storage.
Move `ProfileSource` definition to the .cc file. Change the
`ProfileLoadStatus` to an `enum class` and remove unused
"WouldOverwriteData". Fix an error message.

Test: ProfileCompilationInfoTest.UpdateProfileKeyOkWithAnnotation
Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing --speed-profile
Bug: 148067697
Change-Id: I4e157c76b52e9c9c15f3ba34816f6b606ac4529c
diff --git a/libprofile/profile/profile_compilation_info.cc b/libprofile/profile/profile_compilation_info.cc
index adce530..7b159d8 100644
--- a/libprofile/profile/profile_compilation_info.cc
+++ b/libprofile/profile/profile_compilation_info.cc
@@ -175,7 +175,63 @@
 
 }  // anonymous namespace
 
-struct ProfileCompilationInfo::SafeBuffer {
+/**
+ * Encapsulate the source of profile data for loading.
+ * The source can be either a plain file or a zip file.
+ * For zip files, the profile entry will be extracted to
+ * the memory map.
+ */
+class ProfileCompilationInfo::ProfileSource {
+ public:
+  /**
+   * Create a profile source for the given fd. The ownership of the fd
+   * remains to the caller; as this class will not attempt to close it at any
+   * point.
+   */
+  static ProfileSource* Create(int32_t fd) {
+    DCHECK_GT(fd, -1);
+    return new ProfileSource(fd, MemMap::Invalid());
+  }
+
+  /**
+   * Create a profile source backed by a memory map. The map can be null in
+   * which case it will the treated as an empty source.
+   */
+  static ProfileSource* Create(MemMap&& mem_map) {
+    return new ProfileSource(/*fd*/ -1, std::move(mem_map));
+  }
+
+  /**
+   * Read bytes from this source.
+   * Reading will advance the current source position so subsequent
+   * invocations will read from the las position.
+   */
+  ProfileLoadStatus Read(uint8_t* buffer,
+                         size_t byte_count,
+                         const std::string& debug_stage,
+                         std::string* error);
+
+  /** Return true if the source has 0 data. */
+  bool HasEmptyContent() const;
+
+  /** Return true if all the information from this source has been read. */
+  bool HasConsumedAllData() const;
+
+ private:
+  ProfileSource(int32_t fd, MemMap&& mem_map)
+      : fd_(fd), mem_map_(std::move(mem_map)), mem_map_cur_(0) {}
+
+  bool IsMemMap() const {
+    return fd_ == -1;
+  }
+
+  int32_t fd_;  // The fd is not owned by this class.
+  MemMap mem_map_;
+  size_t mem_map_cur_;  // Current position in the map to read from.
+};
+
+// A helper structure to make sure we don't read past our buffers in the loops.
+class ProfileCompilationInfo::SafeBuffer {
  public:
   explicit SafeBuffer(size_t size)
       : storage_(new uint8_t[size]),
@@ -251,7 +307,8 @@
     : default_arena_pool_(),
       allocator_(custom_arena_pool),
       info_(allocator_.Adapter(kArenaAllocProfile)),
-      profile_key_map_(std::less<const std::string>(), allocator_.Adapter(kArenaAllocProfile)) {
+      profile_key_map_(std::less<const std::string_view>(),
+                       allocator_.Adapter(kArenaAllocProfile)) {
   memcpy(version_,
          for_boot_image ? kProfileVersionForBootImage : kProfileVersion,
          kProfileVersionSize);
@@ -268,7 +325,6 @@
 
 ProfileCompilationInfo::~ProfileCompilationInfo() {
   VLOG(profiler) << Dumpable<MemStats>(allocator_.GetMemStats());
-  ClearData();
 }
 
 void ProfileCompilationInfo::DexPcData::AddClass(uint16_t dex_profile_idx,
@@ -389,7 +445,7 @@
   int fd = profile_file->Fd();
 
   ProfileLoadStatus status = LoadInternal(fd, &error);
-  if (status == kProfileLoadSuccess) {
+  if (status == ProfileLoadStatus::kSuccess) {
     return true;
   }
 
@@ -424,12 +480,13 @@
   int fd = profile_file->Fd();
 
   ProfileLoadStatus status = LoadInternal(fd, &error);
-  if (status == kProfileLoadSuccess) {
+  if (status == ProfileLoadStatus::kSuccess) {
     return true;
   }
 
   if (clear_if_invalid &&
-      ((status == kProfileLoadVersionMismatch) || (status == kProfileLoadBadData))) {
+      ((status == ProfileLoadStatus::kVersionMismatch) ||
+       (status == ProfileLoadStatus::kBadData))) {
     LOG(WARNING) << "Clearing bad or obsolete profile data from file "
                  << filename << ": " << error;
     if (profile_file->ClearContent()) {
@@ -559,7 +616,7 @@
   WriteProfileIndex(&buffer, static_cast<ProfileIndexType>(info_.size()));
 
   uint32_t required_capacity = 0;
-  for (const DexFileData* dex_data_ptr : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data_ptr : info_) {
     const DexFileData& dex_data = *dex_data_ptr;
     uint32_t methods_region_size = GetMethodsRegionSize(dex_data);
     required_capacity += kLineHeaderSize +
@@ -590,7 +647,7 @@
   // Dex files must be written in the order of their profile index. This
   // avoids writing the index in the output file and simplifies the parsing logic.
   // Write profile line headers.
-  for (const DexFileData* dex_data_ptr : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data_ptr : info_) {
     const DexFileData& dex_data = *dex_data_ptr;
 
     if (dex_data.profile_key.size() >= kMaxDexFileKeyLength) {
@@ -612,7 +669,7 @@
     AddStringToBuffer(&buffer, dex_data.profile_key);
   }
 
-  for (const DexFileData* dex_data_ptr : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data_ptr : info_) {
     const DexFileData& dex_data = *dex_data_ptr;
 
     // Note that we allow dex files without any methods or classes, so that
@@ -763,31 +820,37 @@
     const std::string& profile_key,
     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() > 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 file. Something went wrong";
+  DCHECK_EQ(profile_key_map_.size(), info_.size());
+  auto profile_index_it = profile_key_map_.lower_bound(profile_key);
+  if (profile_index_it == profile_key_map_.end() || profile_index_it->first != profile_key) {
+    // We did not find the key. Create a new DexFileData if we did not reach the limit.
+    DCHECK_LE(profile_key_map_.size(), MaxProfileIndex());
+    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 file. Something went wrong";
+      }
+      return nullptr;
     }
-    profile_key_map_.erase(profile_key);
-    return nullptr;
-  }
-
-  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(
+    ProfileIndexType new_profile_index = dchecked_integral_cast<ProfileIndexType>(info_.size());
+    std::unique_ptr<DexFileData> dex_file_data(new (&allocator_) DexFileData(
         &allocator_,
         profile_key,
         checksum,
-        profile_index,
+        new_profile_index,
         num_method_ids,
-        IsForBootImage());
-    info_.push_back(dex_file_data);
+        IsForBootImage()));
+    // Record the new data in `profile_key_map_` and `info_`.
+    std::string_view new_key(dex_file_data->profile_key);
+    profile_index_it = profile_key_map_.PutBefore(profile_index_it, new_key, new_profile_index);
+    info_.push_back(std::move(dex_file_data));
+    DCHECK_EQ(profile_key_map_.size(), info_.size());
   }
-  DexFileData* result = info_[profile_index];
+
+  ProfileIndexType profile_index = profile_index_it->second;
+  DexFileData* result = info_[profile_index].get();
 
   // Check that the checksum matches.
   // This may different if for example the dex file was updated and we had a record of the old one.
@@ -822,7 +885,7 @@
   }
 
   ProfileIndexType profile_index = profile_index_it->second;
-  const DexFileData* result = info_[profile_index];
+  const DexFileData* result = info_[profile_index].get();
   if (verify_checksum && !ChecksumMatch(result->checksum, checksum)) {
     return nullptr;
   }
@@ -836,12 +899,12 @@
       const ProfileSampleAnnotation& annotation) const {
   if (annotation == ProfileSampleAnnotation::kNone) {
     std::string_view profile_key = GetProfileDexFileBaseKeyView(dex_file->GetLocation());
-    for (const DexFileData* dex_data : info_) {
+    for (const std::unique_ptr<DexFileData>& dex_data : info_) {
       if (profile_key == GetBaseKeyViewFromAugmentedKey(dex_data->profile_key)) {
         if (!ChecksumMatch(dex_data->checksum, dex_file->GetLocationChecksum())) {
           return nullptr;
         }
-        return dex_data;
+        return dex_data.get();
       }
     }
   } else {
@@ -856,10 +919,10 @@
     const DexFile* dex_file,
     /*out*/ std::vector<const ProfileCompilationInfo::DexFileData*>* result) const {
   std::string_view profile_key = GetProfileDexFileBaseKeyView(dex_file->GetLocation());
-  for (const DexFileData* dex_data : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data : info_) {
     if (profile_key == GetBaseKeyViewFromAugmentedKey(dex_data->profile_key)) {
       if (ChecksumMatch(dex_data->checksum, dex_file->GetLocationChecksum())) {
-        result->push_back(dex_data);
+        result->push_back(dex_data.get());
       }
     }
   }
@@ -1068,23 +1131,23 @@
   SafeBuffer safe_buffer_version(kMagicVersionSize);
 
   ProfileLoadStatus status = safe_buffer_version.Fill(source, "ReadProfileHeaderVersion", error);
-  if (status != kProfileLoadSuccess) {
+  if (status != ProfileLoadStatus::kSuccess) {
     return status;
   }
 
   if (!safe_buffer_version.CompareAndAdvance(kProfileMagic, sizeof(kProfileMagic))) {
     *error = "Profile missing magic";
-    return kProfileLoadVersionMismatch;
+    return ProfileLoadStatus::kVersionMismatch;
   }
   if (safe_buffer_version.CountUnreadBytes() < kProfileVersionSize) {
      *error = "Cannot read profile version";
-     return kProfileLoadBadData;
+     return ProfileLoadStatus::kBadData;
   }
   memcpy(version_, safe_buffer_version.GetCurrentPtr(), kProfileVersionSize);
   if ((memcmp(version_, kProfileVersion, kProfileVersionSize) != 0) &&
       (memcmp(version_, kProfileVersionForBootImage, kProfileVersionSize) != 0)) {
     *error = "Profile version mismatch";
-    return kProfileLoadVersionMismatch;
+    return ProfileLoadStatus::kVersionMismatch;
   }
 
   const size_t kProfileHeaderDataSize =
@@ -1094,23 +1157,23 @@
   SafeBuffer safe_buffer_header_data(kProfileHeaderDataSize);
 
   status = safe_buffer_header_data.Fill(source, "ReadProfileHeaderData", error);
-  if (status != kProfileLoadSuccess) {
+  if (status != ProfileLoadStatus::kSuccess) {
     return status;
   }
 
   if (!ReadProfileIndex(safe_buffer_header_data, number_of_dex_files)) {
     *error = "Cannot read the number of dex files";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
   if (!safe_buffer_header_data.ReadUintAndAdvance<uint32_t>(uncompressed_data_size)) {
     *error = "Cannot read the size of uncompressed data";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
   if (!safe_buffer_header_data.ReadUintAndAdvance<uint32_t>(compressed_data_size)) {
     *error = "Cannot read the size of compressed data";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
-  return kProfileLoadSuccess;
+  return ProfileLoadStatus::kSuccess;
 }
 
 bool ProfileCompilationInfo::ReadProfileLineHeaderElements(SafeBuffer& buffer,
@@ -1131,29 +1194,29 @@
     /*out*/std::string* error) {
   if (buffer.CountUnreadBytes() < kLineHeaderSize) {
     *error += "Profile EOF reached prematurely for ReadProfileLineHeader";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   uint16_t profile_key_size;
   if (!ReadProfileLineHeaderElements(buffer, &profile_key_size, line_header, error)) {
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   if (profile_key_size == 0 || profile_key_size > kMaxDexFileKeyLength) {
     *error = "ProfileKey has an invalid size: " +
         std::to_string(static_cast<uint32_t>(profile_key_size));
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   if (buffer.CountUnreadBytes() < profile_key_size) {
     *error += "Profile EOF reached prematurely for ReadProfileHeaderDexLocation";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
   const uint8_t* base_ptr = buffer.GetCurrentPtr();
   line_header->profile_key.assign(
       reinterpret_cast<const char*>(base_ptr), profile_key_size);
   buffer.Advance(profile_key_size);
-  return kProfileLoadSuccess;
+  return ProfileLoadStatus::kSuccess;
 }
 
 ProfileCompilationInfo::ProfileLoadStatus ProfileCompilationInfo::ReadProfileLine(
@@ -1169,30 +1232,30 @@
   if (data == nullptr) {
     *error = "Error when reading profile file line header: checksum mismatch for "
         + line_header.profile_key;
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   if (!ReadMethods(buffer, number_of_dex_files, line_header, dex_profile_index_remap, error)) {
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   if (merge_classes) {
     if (!ReadClasses(buffer, line_header, error)) {
-      return kProfileLoadBadData;
+      return ProfileLoadStatus::kBadData;
     }
   }
 
   // Read method bitmap.
   const size_t bytes = data->bitmap_storage.size();
   if (buffer.CountUnreadBytes() < bytes) {
-    *error += "Profile EOF reached prematurely for ReadProfileHeaderDexLocation";
-    return kProfileLoadBadData;
+    *error += "Profile EOF reached prematurely for method bitmap";
+    return ProfileLoadStatus::kBadData;
   }
   const uint8_t* base_ptr = buffer.GetCurrentPtr();
   std::copy_n(base_ptr, bytes, data->bitmap_storage.data());
   buffer.Advance(bytes);
 
-  return kProfileLoadSuccess;
+  return ProfileLoadStatus::kSuccess;
 }
 
 // TODO(calin): Fix this API. ProfileCompilationInfo::Load should be static and
@@ -1203,7 +1266,7 @@
 
   ProfileLoadStatus status = LoadInternal(fd, &error, merge_classes, filter_fn);
 
-  if (status == kProfileLoadSuccess) {
+  if (status == ProfileLoadStatus::kSuccess) {
     return true;
   } else {
     LOG(WARNING) << "Error when reading profile: " << error;
@@ -1216,7 +1279,7 @@
   for (const DexFile* dex_file : dex_files) {
     key_to_dex_file.emplace(GetProfileDexFileBaseKeyView(dex_file->GetLocation()), dex_file);
   }
-  for (const DexFileData* dex_data : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data : info_) {
     // We need to remove any annotation from the key during verification.
     const auto it = key_to_dex_file.find(GetBaseKeyViewFromAugmentedKey(dex_data->profile_key));
     if (it == key_to_dex_file.end()) {
@@ -1310,13 +1373,13 @@
     /*out*/ std::string* error) {
   if (IsProfileFile(fd)) {
     source->reset(ProfileSource::Create(fd));
-    return kProfileLoadSuccess;
+    return ProfileLoadStatus::kSuccess;
   } else {
     std::unique_ptr<ZipArchive> zip_archive(
         ZipArchive::OpenFromFd(DupCloexec(fd), "profile", error));
     if (zip_archive.get() == nullptr) {
       *error = "Could not open the profile zip archive";
-      return kProfileLoadBadData;
+      return ProfileLoadStatus::kBadData;
     }
     std::unique_ptr<ZipEntry> zip_entry(zip_archive->Find(kDexMetadataProfileEntry, error));
     if (zip_entry == nullptr) {
@@ -1326,11 +1389,11 @@
       LOG(WARNING) << "Could not find entry " << kDexMetadataProfileEntry
           << " in the zip archive. Creating an empty profile.";
       source->reset(ProfileSource::Create(MemMap::Invalid()));
-      return kProfileLoadSuccess;
+      return ProfileLoadStatus::kSuccess;
     }
     if (zip_entry->GetUncompressedLength() == 0) {
       *error = "Empty profile entry in the zip archive.";
-      return kProfileLoadBadData;
+      return ProfileLoadStatus::kBadData;
     }
 
     // TODO(calin) pass along file names to assist with debugging.
@@ -1339,9 +1402,9 @@
 
     if (map.IsValid()) {
       source->reset(ProfileSource::Create(std::move(map)));
-      return kProfileLoadSuccess;
+      return ProfileLoadStatus::kSuccess;
     } else {
-      return kProfileLoadBadData;
+      return ProfileLoadStatus::kBadData;
     }
   }
 }
@@ -1353,7 +1416,7 @@
     std::string* error) {
   if (IsMemMap()) {
     if (mem_map_cur_ + byte_count > mem_map_.Size()) {
-      return kProfileLoadBadData;
+      return ProfileLoadStatus::kBadData;
     }
     for (size_t i = 0; i < byte_count; i++) {
       buffer[i] = *(mem_map_.Begin() + mem_map_cur_);
@@ -1364,16 +1427,16 @@
       int bytes_read = TEMP_FAILURE_RETRY(read(fd_, buffer, byte_count));;
       if (bytes_read == 0) {
         *error += "Profile EOF reached prematurely for " + debug_stage;
-        return kProfileLoadBadData;
+        return ProfileLoadStatus::kBadData;
       } else if (bytes_read < 0) {
         *error += "Profile IO error for " + debug_stage + strerror(errno);
-        return kProfileLoadIOError;
+        return ProfileLoadStatus::kIOError;
       }
       byte_count -= bytes_read;
       buffer += bytes_read;
     }
   }
-  return kProfileLoadSuccess;
+  return ProfileLoadStatus::kSuccess;
 }
 
 bool ProfileCompilationInfo::ProfileSource::HasConsumedAllData() const {
@@ -1405,7 +1468,7 @@
 
   std::unique_ptr<ProfileSource> source;
   ProfileLoadStatus status = OpenSource(fd, &source, error);
-  if (status != kProfileLoadSuccess) {
+  if (status != ProfileLoadStatus::kSuccess) {
     return status;
   }
 
@@ -1413,7 +1476,7 @@
   // Profiles may be created by ActivityManager or installd before we manage to
   // process them in the runtime or profman.
   if (source->HasEmptyContent()) {
-    return kProfileLoadSuccess;
+    return ProfileLoadStatus::kSuccess;
   }
 
   // Read profile header: magic + version + number_of_dex_files.
@@ -1426,7 +1489,7 @@
                              &compressed_data_size,
                              error);
 
-  if (status != kProfileLoadSuccess) {
+  if (status != ProfileLoadStatus::kSuccess) {
     return status;
   }
   // Allow large profiles for non target builds for the case where we are merging many profiles
@@ -1435,7 +1498,7 @@
     LOG(ERROR) << "Profile data size exceeds "
                << GetSizeErrorThresholdBytes()
                << " bytes. It has " << uncompressed_data_size << " bytes.";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
   if (uncompressed_data_size > GetSizeWarningThresholdBytes()) {
     LOG(WARNING) << "Profile data size exceeds "
@@ -1445,14 +1508,14 @@
 
   std::unique_ptr<uint8_t[]> compressed_data(new uint8_t[compressed_data_size]);
   status = source->Read(compressed_data.get(), compressed_data_size, "ReadContent", error);
-  if (status != kProfileLoadSuccess) {
+  if (status != ProfileLoadStatus::kSuccess) {
     *error += "Unable to read compressed profile data";
     return status;
   }
 
   if (!source->HasConsumedAllData()) {
     *error += "Unexpected data in the profile file.";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   SafeBuffer uncompressed_data(uncompressed_data_size);
@@ -1463,7 +1526,7 @@
 
   if (ret != Z_STREAM_END) {
     *error += "Error reading uncompressed profile data";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   std::vector<ProfileLineHeader> profile_line_headers;
@@ -1473,7 +1536,7 @@
 
     // First, read the line header to get the amount of data we need to read.
     status = ReadProfileLineHeader(uncompressed_data, &line_header, error);
-    if (status != kProfileLoadSuccess) {
+    if (status != ProfileLoadStatus::kSuccess) {
       return status;
     }
     profile_line_headers.push_back(line_header);
@@ -1481,7 +1544,7 @@
 
   SafeMap<ProfileIndexType, ProfileIndexType> dex_profile_index_remap;
   if (!RemapProfileIndex(profile_line_headers, filter_fn, &dex_profile_index_remap)) {
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   }
 
   for (ProfileIndexType k = 0; k < number_of_dex_files; k++) {
@@ -1501,7 +1564,7 @@
                                dex_profile_index_remap,
                                merge_classes,
                                error);
-      if (status != kProfileLoadSuccess) {
+      if (status != ProfileLoadStatus::kSuccess) {
         return status;
       }
     }
@@ -1511,9 +1574,9 @@
   if (uncompressed_data.CountUnreadBytes() > 0) {
     *error = "Unexpected content in the profile file: " +
         std::to_string(uncompressed_data.CountUnreadBytes()) + " extra bytes";
-    return kProfileLoadBadData;
+    return ProfileLoadStatus::kBadData;
   } else {
-    return kProfileLoadSuccess;
+    return ProfileLoadStatus::kSuccess;
   }
 }
 
@@ -1567,7 +1630,7 @@
   // the current profile info.
   // Note that the number of elements should be very small, so this should not
   // be a performance issue.
-  for (const DexFileData* other_dex_data : other.info_) {
+  for (const std::unique_ptr<DexFileData>& other_dex_data : other.info_) {
     // verify_checksum is false because we want to differentiate between a missing dex data and
     // a mismatched checksum.
     const DexFileData* dex_data = FindDexData(other_dex_data->profile_key,
@@ -1591,7 +1654,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<ProfileIndexType, ProfileIndexType> dex_profile_index_remap;
-  for (const DexFileData* other_dex_data : other.info_) {
+  for (const std::unique_ptr<DexFileData>& other_dex_data : other.info_) {
     const DexFileData* dex_data = GetOrAddDexFileData(other_dex_data->profile_key,
                                                       other_dex_data->checksum,
                                                       other_dex_data->num_method_ids);
@@ -1602,7 +1665,7 @@
   }
 
   // Merge the actual profile data.
-  for (const DexFileData* other_dex_data : other.info_) {
+  for (const std::unique_ptr<DexFileData>& other_dex_data : other.info_) {
     DexFileData* dex_data = const_cast<DexFileData*>(FindDexData(other_dex_data->profile_key,
                                                                  other_dex_data->checksum));
     DCHECK(dex_data != nullptr);
@@ -1663,7 +1726,7 @@
 
 uint32_t ProfileCompilationInfo::GetNumberOfMethods() const {
   uint32_t total = 0;
-  for (const DexFileData* dex_data : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data : info_) {
     total += dex_data->method_map.size();
   }
   return total;
@@ -1671,7 +1734,7 @@
 
 uint32_t ProfileCompilationInfo::GetNumberOfResolvedClasses() const {
   uint32_t total = 0;
-  for (const DexFileData* dex_data : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data : info_) {
     total += dex_data->class_set.size();
   }
   return total;
@@ -1697,7 +1760,7 @@
 
   const std::string kFirstDexFileKeySubstitute = "!classes.dex";
 
-  for (const DexFileData* dex_data : info_) {
+  for (const std::unique_ptr<DexFileData>& dex_data : info_) {
     os << "\n";
     if (print_full_dex_location) {
       os << dex_data->profile_key;
@@ -2097,7 +2160,7 @@
 bool ProfileCompilationInfo::UpdateProfileKeys(
       const std::vector<std::unique_ptr<const DexFile>>& dex_files) {
   for (const std::unique_ptr<const DexFile>& dex_file : dex_files) {
-    for (DexFileData* dex_data : info_) {
+    for (const std::unique_ptr<DexFileData>& dex_data : info_) {
       if (dex_data->checksum == dex_file->GetLocationChecksum()
           && dex_data->num_method_ids == dex_file->NumMethodIds()) {
         std::string new_profile_key = GetProfileDexFileBaseKey(dex_file->GetLocation());
@@ -2112,9 +2175,8 @@
           profile_key_map_.erase(dex_data->profile_key);
           // Retain the annotation (if any) during the renaming by re-attaching the info
           // form the old key.
-          profile_key_map_.Put(MigrateAnnotationInfo(new_profile_key, dex_data->profile_key),
-                               dex_data->profile_index);
-          dex_data->profile_key = new_profile_key;
+          dex_data->profile_key = MigrateAnnotationInfo(new_profile_key, dex_data->profile_key);
+          profile_key_map_.Put(dex_data->profile_key, dex_data->profile_index);
         }
       }
     }
@@ -2129,11 +2191,8 @@
 }
 
 void ProfileCompilationInfo::ClearData() {
-  for (DexFileData* data : info_) {
-    delete data;
-  }
-  info_.clear();
   profile_key_map_.clear();
+  info_.clear();
 }
 
 void ProfileCompilationInfo::ClearDataAndAdjustVersion(bool for_boot_image) {
diff --git a/libprofile/profile/profile_compilation_info.h b/libprofile/profile/profile_compilation_info.h
index 56ca2fd..99794d1 100644
--- a/libprofile/profile/profile_compilation_info.h
+++ b/libprofile/profile/profile_compilation_info.h
@@ -20,6 +20,7 @@
 #include <array>
 #include <list>
 #include <set>
+#include <string_view>
 #include <vector>
 
 #include "base/arena_containers.h"
@@ -445,7 +446,7 @@
     static_assert(std::is_same_v<typename Container::value_type, const DexFile*> ||
                   std::is_same_v<typename Container::value_type, std::unique_ptr<const DexFile>>);
     DCHECK_LE(profile_index, info_.size());
-    const DexFileData* dex_file_data = info_[profile_index];
+    const DexFileData* dex_file_data = info_[profile_index].get();
     DCHECK(dex_file_data != nullptr);
     uint32_t dex_checksum = dex_file_data->checksum;
     std::string_view base_key = GetBaseKeyViewFromAugmentedKey(dex_file_data->profile_key);
@@ -575,20 +576,22 @@
       const std::vector<std::unique_ptr<const DexFile>>& dex_files) const;
 
  private:
-  enum ProfileLoadStatus {
-    kProfileLoadWouldOverwiteData,
-    kProfileLoadIOError,
-    kProfileLoadVersionMismatch,
-    kProfileLoadBadData,
-    kProfileLoadSuccess
+  // Helper classes.
+  class ProfileSource;
+  class SafeBuffer;
+
+  enum class ProfileLoadStatus : uint32_t {
+    kSuccess,
+    kIOError,
+    kVersionMismatch,
+    kBadData,
   };
 
   // Internal representation of the profile information belonging to a dex file.
-  // Note that we could do without profile_key (the key used to encode the dex
-  // file in the profile) and profile_index (the index of the dex file in the
-  // profile) fields in this struct because we can infer them from
-  // profile_key_map_ and info_. However, it makes the profiles logic much
-  // simpler if we have references here as well.
+  // Note that we could do without the profile_index (the index of the dex file
+  // in the profile) field in this struct because we can infer it from
+  // `profile_key_map_` and `info_`. However, it makes the profiles logic much
+  // simpler if we have the profile index here as well.
   struct DexFileData : public DeletableArenaObject<kArenaAllocProfile> {
     DexFileData(ArenaAllocator* allocator,
                 const std::string& key,
@@ -720,61 +723,6 @@
     uint32_t num_method_ids;
   };
 
-  /**
-   * Encapsulate the source of profile data for loading.
-   * The source can be either a plain file or a zip file.
-   * For zip files, the profile entry will be extracted to
-   * the memory map.
-   */
-  class ProfileSource {
-   public:
-    /**
-     * Create a profile source for the given fd. The ownership of the fd
-     * remains to the caller; as this class will not attempt to close it at any
-     * point.
-     */
-    static ProfileSource* Create(int32_t fd) {
-      DCHECK_GT(fd, -1);
-      return new ProfileSource(fd, MemMap::Invalid());
-    }
-
-    /**
-     * Create a profile source backed by a memory map. The map can be null in
-     * which case it will the treated as an empty source.
-     */
-    static ProfileSource* Create(MemMap&& mem_map) {
-      return new ProfileSource(/*fd*/ -1, std::move(mem_map));
-    }
-
-    /**
-     * Read bytes from this source.
-     * Reading will advance the current source position so subsequent
-     * invocations will read from the las position.
-     */
-    ProfileLoadStatus Read(uint8_t* buffer,
-                           size_t byte_count,
-                           const std::string& debug_stage,
-                           std::string* error);
-
-    /** Return true if the source has 0 data. */
-    bool HasEmptyContent() const;
-    /** Return true if all the information from this source has been read. */
-    bool HasConsumedAllData() const;
-
-   private:
-    ProfileSource(int32_t fd, MemMap&& mem_map)
-        : fd_(fd), mem_map_(std::move(mem_map)), mem_map_cur_(0) {}
-
-    bool IsMemMap() const { return fd_ == -1; }
-
-    int32_t fd_;  // The fd is not owned by this class.
-    MemMap mem_map_;
-    size_t mem_map_cur_;  // Current position in the map to read from.
-  };
-
-  // A helper structure to make sure we don't read past our buffers in the loops.
-  struct SafeBuffer;
-
   ProfileLoadStatus OpenSource(int32_t fd,
                                /*out*/ std::unique_ptr<ProfileSource>* source,
                                /*out*/ std::string* error);
@@ -908,12 +856,13 @@
   // Vector containing the actual profile info.
   // The vector index is the profile index of the dex data and
   // matched DexFileData::profile_index.
-  ArenaVector<DexFileData*> info_;
+  ArenaVector<std::unique_ptr<DexFileData>> info_;
 
   // 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, ProfileIndexType> profile_key_map_;
+  // The backing storage for the `string_view` is the associated `DexFileData`.
+  ArenaSafeMap<const std::string_view, ProfileIndexType> profile_key_map_;
 
   // The version of the profile.
   uint8_t version_[kProfileVersionSize];
@@ -1011,7 +960,7 @@
 
 inline ProfileCompilationInfo::DexReferenceDumper ProfileCompilationInfo::DumpDexReference(
     ProfileIndexType profile_index) const {
-  return DexReferenceDumper{info_[profile_index]};
+  return DexReferenceDumper{info_[profile_index].get()};
 }
 
 std::ostream& operator<<(std::ostream& stream, ProfileCompilationInfo::DexReferenceDumper dumper);
diff --git a/libprofile/profile/profile_compilation_info_test.cc b/libprofile/profile/profile_compilation_info_test.cc
index 1e5cdc5..a279717 100644
--- a/libprofile/profile/profile_compilation_info_test.cc
+++ b/libprofile/profile/profile_compilation_info_test.cc
@@ -922,6 +922,32 @@
   }
 }
 
+TEST_F(ProfileCompilationInfoTest, UpdateProfileKeyOkWithAnnotation) {
+  std::vector<std::unique_ptr<const DexFile>> dex_files;
+  dex_files.push_back(std::unique_ptr<const DexFile>(dex1_renamed));
+  dex_files.push_back(std::unique_ptr<const DexFile>(dex2_renamed));
+
+  ProfileCompilationInfo info;
+  ProfileCompilationInfo::ProfileSampleAnnotation annotation("test.package");
+  AddMethod(&info, dex1, /* method_idx= */ 0, Hotness::kFlagHot, annotation);
+  AddMethod(&info, dex2, /* method_idx= */ 0, Hotness::kFlagHot, annotation);
+
+  // Update the profile keys based on the original dex files
+  ASSERT_TRUE(info.UpdateProfileKeys(dex_files));
+
+  // Verify that we find the methods when searched with the original dex files.
+  for (const std::unique_ptr<const DexFile>& dex : dex_files) {
+    ProfileCompilationInfo::MethodHotness loaded_hotness =
+        GetMethod(info, dex.get(), /* method_idx= */ 0, annotation);
+    ASSERT_TRUE(loaded_hotness.IsHot());
+  }
+
+  // Release the ownership as this is held by the test class;
+  for (std::unique_ptr<const DexFile>& dex : dex_files) {
+    UNUSED(dex.release());
+  }
+}
+
 TEST_F(ProfileCompilationInfoTest, UpdateProfileKeyOkButNoUpdate) {
   std::vector<std::unique_ptr<const DexFile>> dex_files;
   dex_files.push_back(std::unique_ptr<const DexFile>(dex1));