Clean up profile AddMethod APIs

Restructure the profile assistant test to use the same APIs as profile
saver and
clean up all the other usages for AddMethodIndex and AddMethod.

Bug: 139884006
Test: m test-art-host

Change-Id: Icf76d6aa05f0f9dcc589182196ca34a0298f2cb7
diff --git a/dexlayout/dexlayout_test.cc b/dexlayout/dexlayout_test.cc
index c3abaa5..b4e023f 100644
--- a/dexlayout/dexlayout_test.cc
+++ b/dexlayout/dexlayout_test.cc
@@ -353,11 +353,8 @@
           flags |= ProfileCompilationInfo::MethodHotness::kFlagStartup;
           ++profile_methods;
         }
-        pfi.AddMethodIndex(static_cast<ProfileCompilationInfo::MethodHotness::Flag>(flags),
-                           dex_file->GetLocation(),
-                           dex_file->GetLocationChecksum(),
-                           /*method_idx=*/i,
-                           dex_file->NumMethodIds());
+        pfi.AddMethod(ProfileMethodInfo(MethodReference(dex_file.get(), /*index=*/ i)),
+                      static_cast<ProfileCompilationInfo::MethodHotness::Flag>(flags));
       }
       // Add every even class too.
       std::set<dex::TypeIndex> classes;
diff --git a/libprofile/profile/profile_compilation_info.cc b/libprofile/profile/profile_compilation_info.cc
index e951f2b..afefa79 100644
--- a/libprofile/profile/profile_compilation_info.cc
+++ b/libprofile/profile/profile_compilation_info.cc
@@ -172,20 +172,6 @@
   }
 }
 
-bool ProfileCompilationInfo::AddMethodIndex(MethodHotness::Flag flags,
-                                            const std::string& dex_location,
-                                            uint32_t checksum,
-                                            uint16_t method_idx,
-                                            uint32_t num_method_ids) {
-  DexFileData* data = GetOrAddDexFileData(GetProfileDexFileKey(dex_location),
-                                          checksum,
-                                          num_method_ids);
-  if (data == nullptr) {
-    return false;
-  }
-  return data->AddMethod(flags, method_idx);
-}
-
 bool ProfileCompilationInfo::AddMethods(const std::vector<ProfileMethodInfo>& methods,
                                         MethodHotness::Flag flags) {
   for (const ProfileMethodInfo& method : methods) {
@@ -655,73 +641,6 @@
   return result;
 }
 
-bool ProfileCompilationInfo::AddMethod(const std::string& dex_location,
-                                       uint32_t dex_checksum,
-                                       uint16_t method_index,
-                                       uint32_t num_method_ids,
-                                       const OfflineProfileMethodInfo& pmi,
-                                       MethodHotness::Flag flags) {
-  DexFileData* const data = GetOrAddDexFileData(GetProfileDexFileKey(dex_location),
-                                                dex_checksum,
-                                                num_method_ids);
-  if (data == nullptr) {
-    // The data is null if there is a mismatch in the checksum or number of method ids.
-    return false;
-  }
-  if (!data->AddMethod(flags, method_index)) {
-    // Happens if the method index is outside the range (i.e. is greater then the number
-    // of methods in the dex file). This should not happen during normal execution,
-    // But tools (e.g. boot image aggregation tools) and tests stress this behaviour.
-    return false;
-  }
-  if ((flags & MethodHotness::kFlagHot) == 0) {
-    // The method is not hot, do not add inline caches.
-    return true;
-  }
-
-  // Add inline caches.
-  InlineCacheMap* inline_cache = data->FindOrAddMethod(method_index);
-  DCHECK(inline_cache != nullptr);
-
-  data->SetMethodHotness(method_index, flags);
-
-  if (pmi.inline_caches == nullptr) {
-    // If we don't have inline caches return success right away.
-    return true;
-  }
-  for (const auto& pmi_inline_cache_it : *pmi.inline_caches) {
-    uint16_t pmi_ic_dex_pc = pmi_inline_cache_it.first;
-    const DexPcData& pmi_ic_dex_pc_data = pmi_inline_cache_it.second;
-    DexPcData* dex_pc_data = FindOrAddDexPc(inline_cache, pmi_ic_dex_pc);
-    if (dex_pc_data->is_missing_types || dex_pc_data->is_megamorphic) {
-      // We are already megamorphic or we are missing types; no point in going forward.
-      continue;
-    }
-
-    if (pmi_ic_dex_pc_data.is_missing_types) {
-      dex_pc_data->SetIsMissingTypes();
-      continue;
-    }
-    if (pmi_ic_dex_pc_data.is_megamorphic) {
-      dex_pc_data->SetIsMegamorphic();
-      continue;
-    }
-
-    for (const ClassReference& class_ref : pmi_ic_dex_pc_data.classes) {
-      const DexReference& dex_ref = pmi.dex_references[class_ref.dex_profile_index];
-      DexFileData* class_dex_data = GetOrAddDexFileData(
-          dex_ref.profile_key,
-          dex_ref.dex_checksum,
-          dex_ref.num_method_ids);
-      if (class_dex_data == nullptr) {  // checksum mismatch
-        return false;
-      }
-      dex_pc_data->AddClass(class_dex_data->profile_index, class_ref.type_index);
-    }
-  }
-  return true;
-}
-
 bool ProfileCompilationInfo::AddMethod(const ProfileMethodInfo& pmi, MethodHotness::Flag flags) {
   DexFileData* const data = GetOrAddDexFileData(pmi.ref.dex_file);
   if (data == nullptr) {  // checksum mismatch
@@ -1855,6 +1774,7 @@
     std::string dex_location = DexFileLoader::GetMultiDexLocation(i, base_dex_location.c_str());
     std::string profile_key = info.GetProfileDexFileKey(dex_location);
 
+    DexFileData* const data = info.GetOrAddDexFileData(profile_key, /*checksum=*/ 0, max_method);
     for (uint16_t m = 0; m < number_of_methods; m++) {
       uint16_t method_idx = rand() % max_method;
       if (m < (number_of_methods / kFavorSplit)) {
@@ -1863,11 +1783,7 @@
       // Alternate between startup and post startup.
       uint32_t flags = MethodHotness::kFlagHot;
       flags |= ((m & 1) != 0) ? MethodHotness::kFlagPostStartup : MethodHotness::kFlagStartup;
-      info.AddMethodIndex(static_cast<MethodHotness::Flag>(flags),
-                          profile_key,
-                          /*checksum=*/ 0,
-                          method_idx,
-                          max_method);
+      data->AddMethod(static_cast<MethodHotness::Flag>(flags), method_idx);
     }
 
     for (uint16_t c = 0; c < number_of_classes; c++) {
@@ -1875,7 +1791,7 @@
       if (c < (number_of_classes / kFavorSplit)) {
         type_idx %= kFavorFirstN;
       }
-      info.AddClassIndex(profile_key, 0, dex::TypeIndex(type_idx), max_method);
+      data->class_set.insert(dex::TypeIndex(type_idx));
     }
   }
   return info.Save(fd);
@@ -1909,12 +1825,12 @@
 
     uint32_t number_of_classes = dex_file->NumClassDefs();
     uint32_t classes_required_in_profile = (number_of_classes * class_percentage) / 100;
+
+    DexFileData* const data = info.GetOrAddDexFileData(
+          profile_key, checksum, dex_file->NumMethodIds());
     for (uint32_t class_index : create_shuffled_range(classes_required_in_profile,
                                                       number_of_classes)) {
-      info.AddClassIndex(profile_key,
-                         checksum,
-                         dex_file->GetClassDef(class_index).class_idx_,
-                         dex_file->NumMethodIds());
+      data->class_set.insert(dex_file->GetClassDef(class_index).class_idx_);
     }
 
     uint32_t number_of_methods = dex_file->NumMethodIds();
@@ -1926,8 +1842,7 @@
       flags |= ((method_index & 1) != 0)
                    ? MethodHotness::kFlagPostStartup
                    : MethodHotness::kFlagStartup;
-      info.AddMethod(ProfileMethodInfo(MethodReference(dex_file.get(), method_index)),
-                     static_cast<MethodHotness::Flag>(flags));
+      data->AddMethod(static_cast<MethodHotness::Flag>(flags), method_index);
     }
   }
   return info.Save(fd);
diff --git a/libprofile/profile/profile_compilation_info.h b/libprofile/profile/profile_compilation_info.h
index d766274..125f6e8 100644
--- a/libprofile/profile/profile_compilation_info.h
+++ b/libprofile/profile/profile_compilation_info.h
@@ -285,14 +285,6 @@
     return true;
   }
 
-  // Add a method index to the profile (without inline caches). The method flags determine if it is
-  // hot, startup, or post startup, or a combination of the previous.
-  bool AddMethodIndex(MethodHotness::Flag flags,
-                      const std::string& dex_location,
-                      uint32_t checksum,
-                      uint16_t method_idx,
-                      uint32_t num_method_ids);
-
   // Add a method to the profile using its online representation (containing runtime structures).
   bool AddMethod(const ProfileMethodInfo& pmi, MethodHotness::Flag flags);
 
@@ -591,15 +583,6 @@
                                dex_file->NumMethodIds());
   }
 
-  // Add a method to the profile using its offline representation.
-  // This is mostly used to facilitate testing.
-  bool AddMethod(const std::string& dex_location,
-                 uint32_t dex_checksum,
-                 uint16_t method_index,
-                 uint32_t num_method_ids,
-                 const OfflineProfileMethodInfo& pmi,
-                 MethodHotness::Flag flags);
-
   // Add a class index to the profile.
   bool AddClassIndex(const std::string& profile_key,
                      uint32_t checksum,
diff --git a/profman/profile_assistant_test.cc b/profman/profile_assistant_test.cc
index b603994..e00cc75 100644
--- a/profman/profile_assistant_test.cc
+++ b/profman/profile_assistant_test.cc
@@ -35,78 +35,69 @@
 
 using Hotness = ProfileCompilationInfo::MethodHotness;
 using TypeReferenceSet = std::set<TypeReference, TypeReferenceValueComparator>;
+using ProfileInlineCache = ProfileMethodInfo::ProfileInlineCache;
 
-static constexpr size_t kMaxMethodIds = 65535;
-
+// TODO(calin): These tests share a lot with the ProfileCompilationInfo tests.
+// we should introduce a better abstraction to extract the common parts.
 class ProfileAssistantTest : public CommonRuntimeTest {
  public:
   void PostRuntimeCreate() override {
     allocator_.reset(new ArenaAllocator(Runtime::Current()->GetArenaPool()));
+
+    dex1 = fake_dex_storage.AddFakeDex("location1", /* checksum= */ 1, /* num_method_ids= */ 10001);
+    dex2 = fake_dex_storage.AddFakeDex("location2", /* checksum= */ 2, /* num_method_ids= */ 10002);
+    dex3 = fake_dex_storage.AddFakeDex("location3", /* checksum= */ 3, /* num_method_ids= */ 10003);
+    dex4 = fake_dex_storage.AddFakeDex("location4", /* checksum= */ 4, /* num_method_ids= */ 10004);
+
+    dex1_checksum_missmatch = fake_dex_storage.AddFakeDex(
+        "location1", /* checksum= */ 12, /* num_method_ids= */ 10001);
   }
 
  protected:
-  void SetupProfile(const std::string& id,
-                    uint32_t checksum,
+  bool AddMethod(ProfileCompilationInfo* info,
+                const DexFile* dex,
+                uint16_t method_idx,
+                const std::vector<ProfileInlineCache>& inline_caches,
+                Hotness::Flag flags) {
+    return info->AddMethod(
+        ProfileMethodInfo(MethodReference(dex, method_idx), inline_caches), flags);
+  }
+
+  bool AddMethod(ProfileCompilationInfo* info,
+                 const DexFile* dex,
+                 uint16_t method_idx,
+                 Hotness::Flag flags) {
+    return info->AddMethod(ProfileMethodInfo(MethodReference(dex, method_idx)),
+                           flags);
+  }
+
+  void SetupProfile(const DexFile* dex_file1,
+                    const DexFile* dex_file2,
                     uint16_t number_of_methods,
                     uint16_t number_of_classes,
                     const ScratchFile& profile,
                     ProfileCompilationInfo* info,
                     uint16_t start_method_index = 0,
                     bool reverse_dex_write_order = false) {
-    std::string dex_location1 = "location1" + id;
-    uint32_t dex_location_checksum1 = checksum;
-    std::string dex_location2 = "location2" + id;
-    uint32_t dex_location_checksum2 = 10 * checksum;
-    SetupProfile(dex_location1,
-                 dex_location_checksum1,
-                 dex_location2,
-                 dex_location_checksum2,
-                 number_of_methods,
-                 number_of_classes,
-                 profile,
-                 info,
-                 start_method_index,
-                 reverse_dex_write_order);
-  }
-
-  void SetupProfile(const std::string& dex_location1,
-                    uint32_t dex_location_checksum1,
-                    const std::string& dex_location2,
-                    uint32_t dex_location_checksum2,
-                    uint16_t number_of_methods,
-                    uint16_t number_of_classes,
-                    const ScratchFile& profile,
-                    ProfileCompilationInfo* info,
-                    uint16_t start_method_index = 0,
-                    bool reverse_dex_write_order = false,
-                    uint32_t number_of_methods1 = kMaxMethodIds,
-                    uint32_t number_of_methods2 = kMaxMethodIds) {
     for (uint16_t i = start_method_index; i < start_method_index + number_of_methods; i++) {
       // reverse_dex_write_order controls the order in which the dex files will be added to
       // the profile and thus written to disk.
-      ProfileCompilationInfo::OfflineProfileMethodInfo pmi =
-          GetOfflineProfileMethodInfo(dex_location1, dex_location_checksum1,
-                                      dex_location2, dex_location_checksum2,
-                                      number_of_methods1, number_of_methods2);
+      std::vector<ProfileInlineCache> inline_caches = GetTestInlineCaches(dex_file1 , dex_file2, dex3);
       Hotness::Flag flags = static_cast<Hotness::Flag>(
           Hotness::kFlagHot | Hotness::kFlagPostStartup);
       if (reverse_dex_write_order) {
-        ASSERT_TRUE(info->AddMethod(
-            dex_location2, dex_location_checksum2, i, number_of_methods2, pmi, flags));
-        ASSERT_TRUE(info->AddMethod(
-            dex_location1, dex_location_checksum1, i, number_of_methods1, pmi, flags));
+        ASSERT_TRUE(AddMethod(info, dex_file2, i, inline_caches, flags));
+        ASSERT_TRUE(AddMethod(info, dex_file1, i, inline_caches, flags));
       } else {
-        ASSERT_TRUE(info->AddMethod(
-            dex_location1, dex_location_checksum1, i, number_of_methods1, pmi, flags));
-        ASSERT_TRUE(info->AddMethod(
-            dex_location2, dex_location_checksum2, i, number_of_methods2, pmi, flags));
+        ASSERT_TRUE(AddMethod(info, dex_file1, i, inline_caches, flags));
+        ASSERT_TRUE(AddMethod(info, dex_file2, i, inline_caches, flags));
       }
     }
     for (uint16_t i = 0; i < number_of_classes; i++) {
-      ASSERT_TRUE(info->AddClassIndex(info->GetProfileDexFileKey(dex_location1),
-                                      dex_location_checksum1,
+      ASSERT_TRUE(info->AddClassIndex(info->GetProfileDexFileKey(dex_file1->GetLocation()),
+                                      dex_file1->GetLocationChecksum(),
                                       dex::TypeIndex(i),
-                                      number_of_methods1));
+                                      dex_file1->NumMethodIds()));
     }
 
     ASSERT_TRUE(info->Save(GetFd(profile)));
@@ -114,77 +105,61 @@
     ASSERT_TRUE(profile.GetFile()->ResetOffset());
   }
 
-  void SetupBasicProfile(const std::string& id,
-                         uint32_t checksum,
-                         uint16_t number_of_methods,
+  void SetupBasicProfile(const DexFile* dex,
                          const std::vector<uint32_t>& hot_methods,
                          const std::vector<uint32_t>& startup_methods,
                          const std::vector<uint32_t>& post_startup_methods,
                          const ScratchFile& profile,
                          ProfileCompilationInfo* info) {
-    std::string dex_location = "location1" + id;
     for (uint32_t idx : hot_methods) {
-      info->AddMethodIndex(Hotness::kFlagHot, dex_location, checksum, idx, number_of_methods);
+      AddMethod(info, dex, idx, Hotness::kFlagHot);
     }
     for (uint32_t idx : startup_methods) {
-      info->AddMethodIndex(Hotness::kFlagStartup, dex_location, checksum, idx, number_of_methods);
+      AddMethod(info, dex, idx, Hotness::kFlagStartup);
     }
     for (uint32_t idx : post_startup_methods) {
-      info->AddMethodIndex(Hotness::kFlagPostStartup,
-                           dex_location,
-                           checksum,
-                           idx,
-                           number_of_methods);
+      AddMethod(info, dex, idx, Hotness::kFlagPostStartup);
     }
     ASSERT_TRUE(info->Save(GetFd(profile)));
     ASSERT_EQ(0, profile.GetFile()->Flush());
     ASSERT_TRUE(profile.GetFile()->ResetOffset());
   }
 
-  // Creates an inline cache which will be destructed at the end of the test.
-  ProfileCompilationInfo::InlineCacheMap* CreateInlineCacheMap() {
-    used_inline_caches.emplace_back(new ProfileCompilationInfo::InlineCacheMap(
-        std::less<uint16_t>(), allocator_->Adapter(kArenaAllocProfile)));
-    return used_inline_caches.back().get();
-  }
-
-  ProfileCompilationInfo::OfflineProfileMethodInfo GetOfflineProfileMethodInfo(
-        const std::string& dex_location1, uint32_t dex_checksum1,
-        const std::string& dex_location2, uint32_t dex_checksum2,
-        uint32_t number_of_methods1 = kMaxMethodIds, uint32_t number_of_methods2 = kMaxMethodIds) {
-    ProfileCompilationInfo::InlineCacheMap* ic_map = CreateInlineCacheMap();
-    ProfileCompilationInfo::OfflineProfileMethodInfo pmi(ic_map);
-    pmi.dex_references.emplace_back(dex_location1, dex_checksum1, number_of_methods1);
-    pmi.dex_references.emplace_back(dex_location2, dex_checksum2, number_of_methods2);
-
+  // The dex1_substitute can be used to replace the default dex1 file.
+  std::vector<ProfileInlineCache> GetTestInlineCaches(
+        const DexFile* dex_file1, const DexFile* dex_file2, const DexFile* dex_file3) {
+    std::vector<ProfileInlineCache> inline_caches;
     // Monomorphic
     for (uint16_t dex_pc = 0; dex_pc < 11; dex_pc++) {
-      ProfileCompilationInfo::DexPcData dex_pc_data(allocator_.get());
-      dex_pc_data.AddClass(0, dex::TypeIndex(0));
-      ic_map->Put(dex_pc, dex_pc_data);
+      std::vector<TypeReference> types = {TypeReference(dex_file1, dex::TypeIndex(0))};
+      inline_caches.push_back(ProfileInlineCache(dex_pc, /* missing_types*/ false, types));
     }
     // Polymorphic
     for (uint16_t dex_pc = 11; dex_pc < 22; dex_pc++) {
-      ProfileCompilationInfo::DexPcData dex_pc_data(allocator_.get());
-      dex_pc_data.AddClass(0, dex::TypeIndex(0));
-      dex_pc_data.AddClass(1, dex::TypeIndex(1));
-
-      ic_map->Put(dex_pc, dex_pc_data);
+      std::vector<TypeReference> types = {
+          TypeReference(dex_file1, dex::TypeIndex(0)),
+          TypeReference(dex_file2, dex::TypeIndex(1)),
+          TypeReference(dex_file3, dex::TypeIndex(2))};
+      inline_caches.push_back(ProfileInlineCache(dex_pc, /* missing_types*/ false, types));
     }
     // Megamorphic
     for (uint16_t dex_pc = 22; dex_pc < 33; dex_pc++) {
-      ProfileCompilationInfo::DexPcData dex_pc_data(allocator_.get());
-      dex_pc_data.SetIsMegamorphic();
-      ic_map->Put(dex_pc, dex_pc_data);
+      // we need 5 types to make the cache megamorphic
+      std::vector<TypeReference> types = {
+          TypeReference(dex_file1, dex::TypeIndex(0)),
+          TypeReference(dex_file1, dex::TypeIndex(1)),
+          TypeReference(dex_file1, dex::TypeIndex(2)),
+          TypeReference(dex_file1, dex::TypeIndex(3)),
+          TypeReference(dex_file1, dex::TypeIndex(4))};
+      inline_caches.push_back(ProfileInlineCache(dex_pc, /* missing_types*/ false, types));
     }
     // Missing types
     for (uint16_t dex_pc = 33; dex_pc < 44; dex_pc++) {
-      ProfileCompilationInfo::DexPcData dex_pc_data(allocator_.get());
-      dex_pc_data.SetIsMissingTypes();
-      ic_map->Put(dex_pc, dex_pc_data);
+      std::vector<TypeReference> types;
+      inline_caches.push_back(ProfileInlineCache(dex_pc, /* missing_types*/ true, types));
     }
 
-    return pmi;
+    return inline_caches;
   }
 
   int GetFd(const ScratchFile& file) const {
@@ -389,11 +364,10 @@
       hot_methods_ref.push_back(i);
     }
     ProfileCompilationInfo info1;
-    uint16_t methods_in_profile = std::max(methods_in_cur_profile, methods_in_ref_profile);
-    SetupBasicProfile("p1", 1, methods_in_profile, hot_methods_cur, empty_vector, empty_vector,
+    SetupBasicProfile(dex1, hot_methods_cur, empty_vector, empty_vector,
         profile,  &info1);
     ProfileCompilationInfo info2;
-    SetupBasicProfile("p1", 1, methods_in_profile, hot_methods_ref, empty_vector, empty_vector,
+    SetupBasicProfile(dex1, hot_methods_ref, empty_vector, empty_vector,
         reference_profile,  &info2);
     return ProcessProfiles(profile_fds, reference_profile_fd);
   }
@@ -407,18 +381,20 @@
     int reference_profile_fd = GetFd(reference_profile);
 
     ProfileCompilationInfo info1;
-    SetupProfile("p1", 1, 0, classes_in_cur_profile, profile,  &info1);
+    SetupProfile(dex1, dex2, 0, classes_in_cur_profile, profile,  &info1);
     ProfileCompilationInfo info2;
-    SetupProfile("p1", 1, 0, classes_in_ref_profile, reference_profile, &info2);
+    SetupProfile(dex1, dex2, 0, classes_in_ref_profile, reference_profile, &info2);
     return ProcessProfiles(profile_fds, reference_profile_fd);
   }
 
   std::unique_ptr<ArenaAllocator> allocator_;
 
-  // Cache of inline caches generated during tests.
-  // This makes it easier to pass data between different utilities and ensure that
-  // caches are destructed at the end of the test.
-  std::vector<std::unique_ptr<ProfileCompilationInfo::InlineCacheMap>> used_inline_caches;
+  const DexFile* dex1;
+  const DexFile* dex2;
+  const DexFile* dex3;
+  const DexFile* dex4;
+  const DexFile* dex1_checksum_missmatch;
+  FakeDexStorage fake_dex_storage;
 };
 
 TEST_F(ProfileAssistantTest, AdviseCompilationEmptyReferences) {
@@ -433,9 +409,9 @@
 
   const uint16_t kNumberOfMethodsToEnableCompilation = 100;
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
+  SetupProfile(dex1, dex2, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
   ProfileCompilationInfo info2;
-  SetupProfile("p2", 2, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
+  SetupProfile(dex3, dex4, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
 
   // We should advise compilation.
   ASSERT_EQ(ProfileAssistant::kCompile,
@@ -466,7 +442,7 @@
 
   const uint16_t kNumberOfClassesToEnableCompilation = 100;
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, 0, kNumberOfClassesToEnableCompilation, profile1, &info1);
+  SetupProfile(dex1, dex2, 0, kNumberOfClassesToEnableCompilation, profile1, &info1);
 
   // We should advise compilation.
   ASSERT_EQ(ProfileAssistant::kCompile,
@@ -497,15 +473,15 @@
   // The new profile info will contain the methods with indices 0-100.
   const uint16_t kNumberOfMethodsToEnableCompilation = 100;
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
+  SetupProfile(dex1, dex2, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
   ProfileCompilationInfo info2;
-  SetupProfile("p2", 2, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
+  SetupProfile(dex3, dex4, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
 
 
   // The reference profile info will contain the methods with indices 50-150.
   const uint16_t kNumberOfMethodsAlreadyCompiled = 100;
   ProfileCompilationInfo reference_info;
-  SetupProfile("p1", 1, kNumberOfMethodsAlreadyCompiled, 0, reference_profile,
+  SetupProfile(dex1, dex2, kNumberOfMethodsAlreadyCompiled, 0, reference_profile,
       &reference_info, kNumberOfMethodsToEnableCompilation / 2);
 
   // We should advise compilation.
@@ -540,9 +516,9 @@
 
   const uint16_t kNumberOfMethodsToSkipCompilation = 24;  // Threshold is 100.
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, kNumberOfMethodsToSkipCompilation, 0, profile1, &info1);
+  SetupProfile(dex1, dex2, kNumberOfMethodsToSkipCompilation, 0, profile1, &info1);
   ProfileCompilationInfo info2;
-  SetupProfile("p2", 2, kNumberOfMethodsToSkipCompilation, 0, profile2, &info2);
+  SetupProfile(dex3, dex4, kNumberOfMethodsToSkipCompilation, 0, profile2, &info2);
 
   // We should not advise compilation.
   ASSERT_EQ(ProfileAssistant::kSkipCompilation,
@@ -616,9 +592,10 @@
   const uint16_t kNumberOfMethodsToEnableCompilation = 100;
   // Assign different hashes for the same dex file. This will make merging of information to fail.
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
+  SetupProfile(dex1, dex2, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
   ProfileCompilationInfo info2;
-  SetupProfile("p1", 2, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
+  SetupProfile(
+      dex1_checksum_missmatch, dex2, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
 
   // We should fail processing.
   ASSERT_EQ(ProfileAssistant::kErrorBadProfiles,
@@ -643,9 +620,10 @@
   const uint16_t kNumberOfMethodsToEnableCompilation = 100;
   // Assign different hashes for the same dex file. This will make merging of information to fail.
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
+  SetupProfile(dex1, dex2, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
   ProfileCompilationInfo reference_info;
-  SetupProfile("p1", 2, kNumberOfMethodsToEnableCompilation, 0, reference_profile, &reference_info);
+  SetupProfile(
+      dex1_checksum_missmatch, dex2, kNumberOfMethodsToEnableCompilation, 0, reference_profile, &reference_info);
 
   // We should not advise compilation.
   ASSERT_TRUE(profile1.GetFile()->ResetOffset());
@@ -1016,7 +994,7 @@
   // The new profile info will contain the methods with indices 0-100.
   const uint16_t kNumberOfMethodsToEnableCompilation = 100;
   ProfileCompilationInfo info1;
-  SetupProfile("p1", 1, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1,
+  SetupProfile(dex1, dex2, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1,
       /*start_method_index=*/0, /*reverse_dex_write_order=*/false);
 
   // The reference profile info will contain the methods with indices 50-150.
@@ -1025,7 +1003,7 @@
   // with a different dex order correctly.
   const uint16_t kNumberOfMethodsAlreadyCompiled = 100;
   ProfileCompilationInfo reference_info;
-  SetupProfile("p1", 1, kNumberOfMethodsAlreadyCompiled, 0, reference_profile,
+  SetupProfile(dex1, dex2, kNumberOfMethodsAlreadyCompiled, 0, reference_profile,
       &reference_info, kNumberOfMethodsToEnableCompilation / 2, /*reverse_dex_write_order=*/true);
 
   // We should advise compilation.
@@ -1133,9 +1111,7 @@
   EXPECT_GT(startup_methods.size(), 0u);
   EXPECT_GT(post_startup_methods.size(), 0u);
   ProfileCompilationInfo info1;
-  SetupBasicProfile("p1",
-                    1,
-                    kNumberOfMethods,
+  SetupBasicProfile(dex1,
                     hot_methods,
                     startup_methods,
                     post_startup_methods,
@@ -1188,17 +1164,15 @@
   // The new profile info will contain the methods with indices 0-100.
   const uint16_t kNumberOfMethodsToEnableCompilation = 100;
   ProfileCompilationInfo info1;
-  SetupProfile(d1.GetLocation(), d1.GetLocationChecksum(), "p1", 1,
-      kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
+  SetupProfile(&d1, dex1, kNumberOfMethodsToEnableCompilation, 0, profile1, &info1);
   ProfileCompilationInfo info2;
-  SetupProfile(d2.GetLocation(), d2.GetLocationChecksum(), "p2", 2,
-      kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
+  SetupProfile(&d2, dex2, kNumberOfMethodsToEnableCompilation, 0, profile2, &info2);
 
 
   // The reference profile info will contain the methods with indices 50-150.
   const uint16_t kNumberOfMethodsAlreadyCompiled = 100;
   ProfileCompilationInfo reference_info;
-  SetupProfile(d1.GetLocation(), d1.GetLocationChecksum(), "p1", 1,
+  SetupProfile(&d1, dex1,
       kNumberOfMethodsAlreadyCompiled, 0, reference_profile,
       &reference_info, kNumberOfMethodsToEnableCompilation / 2);
 
@@ -1267,18 +1241,18 @@
 
   ProfileCompilationInfo info1;
   uint16_t num_methods_to_add = std::min(d1.NumMethodIds(), d2.NumMethodIds());
-  SetupProfile("fake-location1",
-               d1.GetLocationChecksum(),
-               "fake-location2",
-               d2.GetLocationChecksum(),
+
+  FakeDexStorage local_storage;
+  const DexFile* dex_to_be_updated1 = local_storage.AddFakeDex(
+      "fake-location1", d1.GetLocationChecksum(), d1.NumMethodIds());
+  const DexFile* dex_to_be_updated2 = local_storage.AddFakeDex(
+      "fake-location2", d2.GetLocationChecksum(), d2.NumMethodIds());
+  SetupProfile(dex_to_be_updated1,
+               dex_to_be_updated2,
                num_methods_to_add,
                /*number_of_classes=*/ 0,
                profile1,
-               &info1,
-               /*start_method_index=*/ 0,
-               /*reverse_dex_write_order=*/ false,
-               /*number_of_methods1=*/ d1.NumMethodIds(),
-               /*number_of_methods2=*/ d2.NumMethodIds());
+               &info1);
 
   // Run profman and pass the dex file with --apk-fd.
   android::base::unique_fd apk_fd(
diff --git a/runtime/jit/profiling_info_test.cc b/runtime/jit/profiling_info_test.cc
index 556aa13..6b82411 100644
--- a/runtime/jit/profiling_info_test.cc
+++ b/runtime/jit/profiling_info_test.cc
@@ -36,8 +36,6 @@
 
 using Hotness = ProfileCompilationInfo::MethodHotness;
 
-static constexpr size_t kMaxMethodIds = 65535;
-
 class ProfileCompilationInfoTest : public CommonRuntimeTest {
  public:
   void PostRuntimeCreate() override {
@@ -63,34 +61,6 @@
     return methods;
   }
 
-  bool AddMethod(const std::string& dex_location,
-                 uint32_t checksum,
-                 uint16_t method_index,
-                 ProfileCompilationInfo* info) {
-    return info->AddMethodIndex(Hotness::kFlagHot,
-                                dex_location,
-                                checksum,
-                                method_index,
-                                kMaxMethodIds);
-  }
-
-  bool AddMethod(const std::string& dex_location,
-                 uint32_t checksum,
-                 uint16_t method_index,
-                 const ProfileCompilationInfo::OfflineProfileMethodInfo& pmi,
-                 ProfileCompilationInfo* info) {
-    return info->AddMethod(
-        dex_location, checksum, method_index, kMaxMethodIds, pmi, Hotness::kFlagPostStartup);
-  }
-
-  bool AddClass(const std::string& dex_location,
-                uint32_t checksum,
-                dex::TypeIndex type_index,
-                ProfileCompilationInfo* info) {
-    return info->AddClassIndex(
-        info->GetProfileDexFileKey(dex_location), checksum, type_index, kMaxMethodIds);
-  }
-
   uint32_t GetFd(const ScratchFile& file) {
     return static_cast<uint32_t>(file.GetFd());
   }
diff --git a/test/ProfileTestMultiDex/Main.java b/test/ProfileTestMultiDex/Main.java
index a8ced54..978cb2c 100644
--- a/test/ProfileTestMultiDex/Main.java
+++ b/test/ProfileTestMultiDex/Main.java
@@ -67,3 +67,157 @@
 class SubE extends Super {
   int getValue() { return 16; };
 }
+
+// Add a class with lots of methods so we can test profile guided compilation triggers.
+class ZLotsOfMethods {
+  public void m1() {}
+  public void m2() {}
+  public void m3() {}
+  public void m4() {}
+  public void m5() {}
+  public void m6() {}
+  public void m7() {}
+  public void m8() {}
+  public void m9() {}
+  public void m10() {}
+  public void m11() {}
+  public void m12() {}
+  public void m13() {}
+  public void m14() {}
+  public void m15() {}
+  public void m16() {}
+  public void m17() {}
+  public void m18() {}
+  public void m19() {}
+  public void m20() {}
+  public void m21() {}
+  public void m22() {}
+  public void m23() {}
+  public void m24() {}
+  public void m25() {}
+  public void m26() {}
+  public void m27() {}
+  public void m28() {}
+  public void m29() {}
+  public void m30() {}
+  public void m31() {}
+  public void m32() {}
+  public void m33() {}
+  public void m34() {}
+  public void m35() {}
+  public void m36() {}
+  public void m37() {}
+  public void m38() {}
+  public void m39() {}
+  public void m40() {}
+  public void m41() {}
+  public void m42() {}
+  public void m43() {}
+  public void m44() {}
+  public void m45() {}
+  public void m46() {}
+  public void m47() {}
+  public void m48() {}
+  public void m49() {}
+  public void m50() {}
+  public void m51() {}
+  public void m52() {}
+  public void m53() {}
+  public void m54() {}
+  public void m55() {}
+  public void m56() {}
+  public void m57() {}
+  public void m58() {}
+  public void m59() {}
+  public void m60() {}
+  public void m61() {}
+  public void m62() {}
+  public void m63() {}
+  public void m64() {}
+  public void m65() {}
+  public void m66() {}
+  public void m67() {}
+  public void m68() {}
+  public void m69() {}
+  public void m70() {}
+  public void m71() {}
+  public void m72() {}
+  public void m73() {}
+  public void m74() {}
+  public void m75() {}
+  public void m76() {}
+  public void m77() {}
+  public void m78() {}
+  public void m79() {}
+  public void m80() {}
+  public void m81() {}
+  public void m82() {}
+  public void m83() {}
+  public void m84() {}
+  public void m85() {}
+  public void m86() {}
+  public void m87() {}
+  public void m88() {}
+  public void m89() {}
+  public void m90() {}
+  public void m91() {}
+  public void m92() {}
+  public void m93() {}
+  public void m94() {}
+  public void m95() {}
+  public void m96() {}
+  public void m97() {}
+  public void m98() {}
+  public void m99() {}
+  public void m100() {}
+  public void m101() {}
+  public void m102() {}
+  public void m103() {}
+  public void m104() {}
+  public void m105() {}
+  public void m106() {}
+  public void m107() {}
+  public void m108() {}
+  public void m109() {}
+  public void m110() {}
+  public void m111() {}
+  public void m112() {}
+  public void m113() {}
+  public void m114() {}
+  public void m115() {}
+  public void m116() {}
+  public void m117() {}
+  public void m118() {}
+  public void m119() {}
+  public void m120() {}
+  public void m121() {}
+  public void m122() {}
+  public void m123() {}
+  public void m124() {}
+  public void m125() {}
+  public void m126() {}
+  public void m127() {}
+  public void m128() {}
+  public void m129() {}
+  public void m130() {}
+  public void m131() {}
+  public void m132() {}
+  public void m133() {}
+  public void m134() {}
+  public void m135() {}
+  public void m136() {}
+  public void m137() {}
+  public void m138() {}
+  public void m139() {}
+  public void m140() {}
+  public void m141() {}
+  public void m142() {}
+  public void m143() {}
+  public void m144() {}
+  public void m145() {}
+  public void m146() {}
+  public void m147() {}
+  public void m148() {}
+  public void m149() {}
+  public void m150() {}
+}
diff --git a/test/ProfileTestMultiDex/Second.java b/test/ProfileTestMultiDex/Second.java
index 9f5dc66..a2bb8d4 100644
--- a/test/ProfileTestMultiDex/Second.java
+++ b/test/ProfileTestMultiDex/Second.java
@@ -38,3 +38,158 @@
     return Integer.valueOf(i);
   }
 }
+
+// Add a class with lots of methods so we can test profile guided compilation triggers.
+// Start the name with 'Z' so that the class is added at the end of the dex file.
+class ZLotsOfMethodsSecond {
+  public void m1() {}
+  public void m2() {}
+  public void m3() {}
+  public void m4() {}
+  public void m5() {}
+  public void m6() {}
+  public void m7() {}
+  public void m8() {}
+  public void m9() {}
+  public void m10() {}
+  public void m11() {}
+  public void m12() {}
+  public void m13() {}
+  public void m14() {}
+  public void m15() {}
+  public void m16() {}
+  public void m17() {}
+  public void m18() {}
+  public void m19() {}
+  public void m20() {}
+  public void m21() {}
+  public void m22() {}
+  public void m23() {}
+  public void m24() {}
+  public void m25() {}
+  public void m26() {}
+  public void m27() {}
+  public void m28() {}
+  public void m29() {}
+  public void m30() {}
+  public void m31() {}
+  public void m32() {}
+  public void m33() {}
+  public void m34() {}
+  public void m35() {}
+  public void m36() {}
+  public void m37() {}
+  public void m38() {}
+  public void m39() {}
+  public void m40() {}
+  public void m41() {}
+  public void m42() {}
+  public void m43() {}
+  public void m44() {}
+  public void m45() {}
+  public void m46() {}
+  public void m47() {}
+  public void m48() {}
+  public void m49() {}
+  public void m50() {}
+  public void m51() {}
+  public void m52() {}
+  public void m53() {}
+  public void m54() {}
+  public void m55() {}
+  public void m56() {}
+  public void m57() {}
+  public void m58() {}
+  public void m59() {}
+  public void m60() {}
+  public void m61() {}
+  public void m62() {}
+  public void m63() {}
+  public void m64() {}
+  public void m65() {}
+  public void m66() {}
+  public void m67() {}
+  public void m68() {}
+  public void m69() {}
+  public void m70() {}
+  public void m71() {}
+  public void m72() {}
+  public void m73() {}
+  public void m74() {}
+  public void m75() {}
+  public void m76() {}
+  public void m77() {}
+  public void m78() {}
+  public void m79() {}
+  public void m80() {}
+  public void m81() {}
+  public void m82() {}
+  public void m83() {}
+  public void m84() {}
+  public void m85() {}
+  public void m86() {}
+  public void m87() {}
+  public void m88() {}
+  public void m89() {}
+  public void m90() {}
+  public void m91() {}
+  public void m92() {}
+  public void m93() {}
+  public void m94() {}
+  public void m95() {}
+  public void m96() {}
+  public void m97() {}
+  public void m98() {}
+  public void m99() {}
+  public void m100() {}
+  public void m101() {}
+  public void m102() {}
+  public void m103() {}
+  public void m104() {}
+  public void m105() {}
+  public void m106() {}
+  public void m107() {}
+  public void m108() {}
+  public void m109() {}
+  public void m110() {}
+  public void m111() {}
+  public void m112() {}
+  public void m113() {}
+  public void m114() {}
+  public void m115() {}
+  public void m116() {}
+  public void m117() {}
+  public void m118() {}
+  public void m119() {}
+  public void m120() {}
+  public void m121() {}
+  public void m122() {}
+  public void m123() {}
+  public void m124() {}
+  public void m125() {}
+  public void m126() {}
+  public void m127() {}
+  public void m128() {}
+  public void m129() {}
+  public void m130() {}
+  public void m131() {}
+  public void m132() {}
+  public void m133() {}
+  public void m134() {}
+  public void m135() {}
+  public void m136() {}
+  public void m137() {}
+  public void m138() {}
+  public void m139() {}
+  public void m140() {}
+  public void m141() {}
+  public void m142() {}
+  public void m143() {}
+  public void m144() {}
+  public void m145() {}
+  public void m146() {}
+  public void m147() {}
+  public void m148() {}
+  public void m149() {}
+  public void m150() {}
+}
diff --git a/test/ProfileTestMultiDex/main.jpp b/test/ProfileTestMultiDex/main.jpp
index 5e55e96..0644072 100644
--- a/test/ProfileTestMultiDex/main.jpp
+++ b/test/ProfileTestMultiDex/main.jpp
@@ -19,3 +19,6 @@
 SubE:
   @@com.android.jack.annotations.ForceInMainDex
   class SubE
+ZLotsOfMethods:
+@@com.android.jack.annotations.ForceInMainDex
+  class ZLotsOfMethods
diff --git a/test/ProfileTestMultiDex/main.list b/test/ProfileTestMultiDex/main.list
index ec131f0..6ca79d4 100644
--- a/test/ProfileTestMultiDex/main.list
+++ b/test/ProfileTestMultiDex/main.list
@@ -5,3 +5,4 @@
 SubB.class
 SubD.class
 SubE.class
+ZLotsOfMethods.class