dex2oat: Faster retrieval of profile data.

Avoid the slow `ProfileCompilationInfo::GetMethodHotness()`.
Cache profile index to avoid repeating the `DexFileData`
search and use new specialized functions that retrieve only
the required information instead of full `MethodHotness`.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Bug: 181943478
Change-Id: Iaf8486a5e5b12f114c8abb9cfdedf6fc4ed62e20
diff --git a/dex2oat/driver/compiler_driver.cc b/dex2oat/driver/compiler_driver.cc
index 54ca9b1..cc9f9ec 100644
--- a/dex2oat/driver/compiler_driver.cc
+++ b/dex2oat/driver/compiler_driver.cc
@@ -411,6 +411,40 @@
   }
 }
 
+// Checks whether profile guided compilation is enabled and if the method should be compiled
+// according to the profile file.
+static bool ShouldCompileBasedOnProfile(const CompilerOptions& compiler_options,
+                                        ProfileCompilationInfo::ProfileIndexType profile_index,
+                                        MethodReference method_ref) {
+  if (profile_index == ProfileCompilationInfo::MaxProfileIndex()) {
+    // No profile for this dex file. Check if we're actually compiling based on a profile.
+    if (!CompilerFilter::DependsOnProfile(compiler_options.GetCompilerFilter())) {
+      return true;
+    }
+    // Profile-based compilation without profile for this dex file. Do not compile the method.
+    DCHECK(compiler_options.GetProfileCompilationInfo() == nullptr ||
+           compiler_options.GetProfileCompilationInfo()->FindDexFile(*method_ref.dex_file) ==
+               ProfileCompilationInfo::MaxProfileIndex());
+    return false;
+  } else {
+    DCHECK(CompilerFilter::DependsOnProfile(compiler_options.GetCompilerFilter()));
+    const ProfileCompilationInfo* profile_compilation_info =
+        compiler_options.GetProfileCompilationInfo();
+    DCHECK(profile_compilation_info != nullptr);
+
+    // Compile only hot methods, it is the profile saver's job to decide
+    // what startup methods to mark as hot.
+    bool result = profile_compilation_info->IsHotMethod(profile_index, method_ref.index);
+
+    if (kDebugProfileGuidedCompilation) {
+      LOG(INFO) << "[ProfileGuidedCompilation] "
+          << (result ? "Compiled" : "Skipped") << " method:" << method_ref.PrettyMethod(true);
+    }
+
+    return result;
+  }
+}
+
 static void CompileMethodQuick(
     Thread* self,
     CompilerDriver* driver,
@@ -421,8 +455,9 @@
     uint32_t method_idx,
     Handle<mirror::ClassLoader> class_loader,
     const DexFile& dex_file,
-    Handle<mirror::DexCache> dex_cache) {
-  auto quick_fn = [](
+    Handle<mirror::DexCache> dex_cache,
+    ProfileCompilationInfo::ProfileIndexType profile_index) {
+  auto quick_fn = [profile_index](
       Thread* self ATTRIBUTE_UNUSED,
       CompilerDriver* driver,
       const dex::CodeItem* code_item,
@@ -435,12 +470,12 @@
       Handle<mirror::DexCache> dex_cache) {
     DCHECK(driver != nullptr);
     CompiledMethod* compiled_method = nullptr;
-    MethodReference method_ref(&dex_file, method_idx);
 
     if ((access_flags & kAccNative) != 0) {
       // Are we extracting only and have support for generic JNI down calls?
-      if (!driver->GetCompilerOptions().IsJniCompilationEnabled() &&
-          InstructionSetHasGenericJniStub(driver->GetCompilerOptions().GetInstructionSet())) {
+      const CompilerOptions& compiler_options = driver->GetCompilerOptions();
+      if (!compiler_options.IsJniCompilationEnabled() &&
+          InstructionSetHasGenericJniStub(compiler_options.GetInstructionSet())) {
         // Leaving this empty will trigger the generic JNI version
       } else {
         // Query any JNI optimization annotations such as @FastNative or @CriticalNative.
@@ -454,8 +489,10 @@
     } else if ((access_flags & kAccAbstract) != 0) {
       // Abstract methods don't have code.
     } else {
-      const VerificationResults* results = driver->GetCompilerOptions().GetVerificationResults();
+      const CompilerOptions& compiler_options = driver->GetCompilerOptions();
+      const VerificationResults* results = compiler_options.GetVerificationResults();
       DCHECK(results != nullptr);
+      MethodReference method_ref(&dex_file, method_idx);
       const VerifiedMethod* verified_method = results->GetVerifiedMethod(method_ref);
       bool compile =
           // Basic checks, e.g., not <clinit>.
@@ -466,8 +503,8 @@
           !verified_method->HasRuntimeThrow() &&
           (verified_method->GetEncounteredVerificationFailures() &
               verifier::VERIFY_ERROR_LOCKING) == 0 &&
-              // Is eligable for compilation by methods-to-compile filter.
-              driver->ShouldCompileBasedOnProfile(method_ref);
+          // Is eligible for compilation by methods-to-compile filter.
+          ShouldCompileBasedOnProfile(compiler_options, profile_index, method_ref);
 
       if (compile) {
         // NOTE: if compiler declines to compile this method, it will return null.
@@ -479,11 +516,10 @@
                                                          class_loader,
                                                          dex_file,
                                                          dex_cache);
-        ProfileMethodsCheck check_type =
-            driver->GetCompilerOptions().CheckProfiledMethodsCompiled();
+        ProfileMethodsCheck check_type = compiler_options.CheckProfiledMethodsCompiled();
         if (UNLIKELY(check_type != ProfileMethodsCheck::kNone)) {
-          bool violation = driver->ShouldCompileBasedOnProfile(method_ref) &&
-                               (compiled_method == nullptr);
+          DCHECK(ShouldCompileBasedOnProfile(compiler_options, profile_index, method_ref));
+          bool violation = (compiled_method == nullptr);
           if (violation) {
             std::ostringstream oss;
             oss << "Failed to compile "
@@ -544,7 +580,9 @@
 void CompilerDriver::ResolveConstStrings(const std::vector<const DexFile*>& dex_files,
                                          bool only_startup_strings,
                                          TimingLogger* timings) {
-  if (only_startup_strings && GetCompilerOptions().GetProfileCompilationInfo() == nullptr) {
+  const ProfileCompilationInfo* profile_compilation_info =
+      GetCompilerOptions().GetProfileCompilationInfo();
+  if (only_startup_strings && profile_compilation_info == nullptr) {
     // If there is no profile, don't resolve any strings. Resolving all of the strings in the image
     // will cause a bloated app image and slow down startup.
     return;
@@ -559,15 +597,19 @@
     dex_cache.Assign(class_linker->FindDexCache(soa.Self(), *dex_file));
     TimingLogger::ScopedTiming t("Resolve const-string Strings", timings);
 
+    ProfileCompilationInfo::ProfileIndexType profile_index =
+        ProfileCompilationInfo::MaxProfileIndex();
+    if (profile_compilation_info != nullptr) {
+      profile_index = profile_compilation_info->FindDexFile(*dex_file);
+      if (profile_index == ProfileCompilationInfo::MaxProfileIndex()) {
+        // We have a `ProfileCompilationInfo` but no data for this dex file.
+        // The code below would not find any method to process.
+        continue;
+      }
+    }
+
     // TODO: Implement a profile-based filter for the boot image. See b/76145463.
     for (ClassAccessor accessor : dex_file->GetClasses()) {
-      const ProfileCompilationInfo* profile_compilation_info =
-          GetCompilerOptions().GetProfileCompilationInfo();
-
-      const bool is_startup_class =
-          profile_compilation_info != nullptr &&
-          profile_compilation_info->ContainsClass(*dex_file, accessor.GetClassIdx());
-
       // Skip methods that failed to verify since they may contain invalid Dex code.
       if (GetClassStatus(ClassReference(dex_file, accessor.GetClassDefIndex())) <
           ClassStatus::kRetryVerificationAtRuntime) {
@@ -575,15 +617,23 @@
       }
 
       for (const ClassAccessor::Method& method : accessor.GetMethods()) {
-        const bool is_clinit = (method.GetAccessFlags() & kAccConstructor) != 0 &&
-            (method.GetAccessFlags() & kAccStatic) != 0;
-        const bool is_startup_clinit = is_startup_class && is_clinit;
+        if (profile_compilation_info != nullptr) {
+          DCHECK_NE(profile_index, ProfileCompilationInfo::MaxProfileIndex());
+          // There can be at most one class initializer in a class, so we shall not
+          // call `ProfileCompilationInfo::ContainsClass()` more than once per class.
+          constexpr uint32_t kMask = kAccConstructor | kAccStatic;
+          const bool is_startup_clinit =
+              (method.GetAccessFlags() & kMask) == kMask &&
+              profile_compilation_info->ContainsClass(profile_index, accessor.GetClassIdx());
 
-        if (profile_compilation_info != nullptr && !is_startup_clinit) {
-          ProfileCompilationInfo::MethodHotness hotness =
-              profile_compilation_info->GetMethodHotness(method.GetReference());
-          if (only_startup_strings ? !hotness.IsStartup() : !hotness.IsInProfile()) {
-            continue;
+          if (!is_startup_clinit) {
+            uint32_t method_index = method.GetIndex();
+            bool process_method = only_startup_strings
+                ? profile_compilation_info->IsStartupMethod(profile_index, method_index)
+                : profile_compilation_info->IsMethodInProfile(profile_index, method_index);
+            if (!process_method) {
+              continue;
+            }
           }
         }
 
@@ -859,31 +909,6 @@
   }
 }
 
-bool CompilerDriver::ShouldCompileBasedOnProfile(const MethodReference& method_ref) const {
-  // Profile compilation info may be null if no profile is passed.
-  if (!CompilerFilter::DependsOnProfile(compiler_options_->GetCompilerFilter())) {
-    // Use the compiler filter instead of the presence of profile_compilation_info_ since
-    // we may want to have full speed compilation along with profile based layout optimizations.
-    return true;
-  }
-  // If we are using a profile filter but do not have a profile compilation info, compile nothing.
-  const ProfileCompilationInfo* profile_compilation_info =
-      GetCompilerOptions().GetProfileCompilationInfo();
-  if (profile_compilation_info == nullptr) {
-    return false;
-  }
-  // Compile only hot methods, it is the profile saver's job to decide what startup methods to mark
-  // as hot.
-  bool result = profile_compilation_info->GetMethodHotness(method_ref).IsHot();
-
-  if (kDebugProfileGuidedCompilation) {
-    LOG(INFO) << "[ProfileGuidedCompilation] "
-        << (result ? "Compiled" : "Skipped") << " method:" << method_ref.PrettyMethod(true);
-  }
-
-  return result;
-}
-
 class ResolveCatchBlockExceptionsClassVisitor : public ClassVisitor {
  public:
   ResolveCatchBlockExceptionsClassVisitor() : classes_() {}
@@ -2488,8 +2513,14 @@
                                      &dex_file,
                                      dex_files,
                                      thread_pool);
+  const CompilerOptions& compiler_options = driver->GetCompilerOptions();
+  bool have_profile = (compiler_options.GetProfileCompilationInfo() != nullptr);
+  bool use_profile = CompilerFilter::DependsOnProfile(compiler_options.GetCompilerFilter());
+  ProfileCompilationInfo::ProfileIndexType profile_index = (have_profile && use_profile)
+      ? compiler_options.GetProfileCompilationInfo()->FindDexFile(dex_file)
+      : ProfileCompilationInfo::MaxProfileIndex();
 
-  auto compile = [&context, &compile_fn](size_t class_def_index) {
+  auto compile = [&context, &compile_fn, profile_index](size_t class_def_index) {
     const DexFile& dex_file = *context.GetDexFile();
     SCOPED_TRACE << "compile " << dex_file.GetLocation() << "@" << class_def_index;
     ClassLinker* class_linker = context.GetClassLinker();
@@ -2550,7 +2581,8 @@
                  method_idx,
                  class_loader,
                  dex_file,
-                 dex_cache);
+                 dex_cache,
+                 profile_index);
     }
   };
   context.ForAllLambda(0, dex_file.NumClassDefs(), compile, thread_count);
diff --git a/dex2oat/driver/compiler_driver.h b/dex2oat/driver/compiler_driver.h
index 0a3c0fe..bdabb2e 100644
--- a/dex2oat/driver/compiler_driver.h
+++ b/dex2oat/driver/compiler_driver.h
@@ -194,10 +194,6 @@
     return compiled_method_storage_.DedupeEnabled();
   }
 
-  // Checks whether profile guided compilation is enabled and if the method should be compiled
-  // according to the profile file.
-  bool ShouldCompileBasedOnProfile(const MethodReference& method_ref) const;
-
   // Checks whether profile guided verification is enabled and if the method should be verified
   // according to the profile file.
   bool ShouldVerifyClassBasedOnProfile(const DexFile& dex_file, uint16_t class_idx) const;
diff --git a/dex2oat/linker/oat_writer.cc b/dex2oat/linker/oat_writer.cc
index af2db2d..5bd060e 100644
--- a/dex2oat/linker/oat_writer.cc
+++ b/dex2oat/linker/oat_writer.cc
@@ -1047,7 +1047,7 @@
 //
 // See also OrderedMethodVisitor.
 struct OatWriter::OrderedMethodData {
-  ProfileCompilationInfo::MethodHotness method_hotness;
+  uint32_t hotness_bits;
   OatClass* oat_class;
   CompiledMethod* compiled_method;
   MethodReference method_reference;
@@ -1089,49 +1089,13 @@
     }
 
     // Use the profile's method hotness to determine sort order.
-    if (GetMethodHotnessOrder() < other.GetMethodHotnessOrder()) {
+    if (hotness_bits < other.hotness_bits) {
       return true;
     }
 
     // Default: retain the original order.
     return false;
   }
-
- private:
-  // Used to determine relative order for OAT code layout when determining
-  // binning.
-  size_t GetMethodHotnessOrder() const {
-    bool hotness[] = {
-      method_hotness.IsHot(),
-      method_hotness.IsStartup(),
-      method_hotness.IsPostStartup()
-    };
-
-
-    // Note: Bin-to-bin order does not matter. If the kernel does or does not read-ahead
-    // any memory, it only goes into the buffer cache and does not grow the PSS until the first
-    // time that memory is referenced in the process.
-
-    size_t hotness_bits = 0;
-    for (size_t i = 0; i < arraysize(hotness); ++i) {
-      if (hotness[i]) {
-        hotness_bits |= (1 << i);
-      }
-    }
-
-    if (kIsDebugBuild) {
-      // Check for bins that are always-empty given a real profile.
-      if (method_hotness.IsHot() &&
-              !method_hotness.IsStartup() && !method_hotness.IsPostStartup()) {
-        std::string name = method_reference.PrettyMethod();
-        LOG(FATAL) << "Method " << name << " had a Hot method that wasn't marked "
-                   << "either start-up or post-startup. Possible corrupted profile?";
-        // This is not fatal, so only warn.
-      }
-    }
-
-    return hotness_bits;
-  }
 };
 
 // Given a queue of CompiledMethod in some total order,
@@ -1190,7 +1154,23 @@
 class OatWriter::LayoutCodeMethodVisitor : public OatDexMethodVisitor {
  public:
   LayoutCodeMethodVisitor(OatWriter* writer, size_t offset)
-      : OatDexMethodVisitor(writer, offset) {
+      : OatDexMethodVisitor(writer, offset),
+        profile_index_(ProfileCompilationInfo::MaxProfileIndex()),
+        profile_index_dex_file_(nullptr) {
+  }
+
+  bool StartClass(const DexFile* dex_file, size_t class_def_index) override {
+    // Update the cached `profile_index_` if needed. This happens only once per dex file
+    // because we visit all classes in a dex file together, so mark that as `UNLIKELY`.
+    if (UNLIKELY(dex_file != profile_index_dex_file_)) {
+      if (writer_->profile_compilation_info_ != nullptr) {
+        profile_index_ = writer_->profile_compilation_info_->FindDexFile(*dex_file);
+      } else {
+        DCHECK_EQ(profile_index_, ProfileCompilationInfo::MaxProfileIndex());
+      }
+      profile_index_dex_file_ = dex_file;
+    }
+    return OatDexMethodVisitor::StartClass(dex_file, class_def_index);
   }
 
   bool EndClass() override {
@@ -1231,18 +1211,37 @@
         }
       }
 
-      MethodReference method_ref(dex_file_, method.GetIndex());
-
-      // Lookup method hotness from profile, if available.
-      // Otherwise assume a default of none-hotness.
-      ProfileCompilationInfo::MethodHotness method_hotness =
-          writer_->profile_compilation_info_ != nullptr
-              ? writer_->profile_compilation_info_->GetMethodHotness(method_ref)
-              : ProfileCompilationInfo::MethodHotness();
+      // Determine the `hotness_bits`, used to determine relative order
+      // for OAT code layout when determining binning.
+      uint32_t method_index = method.GetIndex();
+      MethodReference method_ref(dex_file_, method_index);
+      uint32_t hotness_bits = 0u;
+      if (profile_index_ != ProfileCompilationInfo::MaxProfileIndex()) {
+        ProfileCompilationInfo* pci = writer_->profile_compilation_info_;
+        DCHECK(pci != nullptr);
+        // Note: Bin-to-bin order does not matter. If the kernel does or does not read-ahead
+        // any memory, it only goes into the buffer cache and does not grow the PSS until the
+        // first time that memory is referenced in the process.
+        constexpr uint32_t kHotBit = 1u;
+        constexpr uint32_t kStartupBit = 2u;
+        constexpr uint32_t kPostStartupBit = 4u;
+        hotness_bits =
+            (pci->IsHotMethod(profile_index_, method_index) ? kHotBit : 0u) |
+            (pci->IsStartupMethod(profile_index_, method_index) ? kStartupBit : 0u) |
+            (pci->IsPostStartupMethod(profile_index_, method_index) ? kPostStartupBit : 0u);
+        if (kIsDebugBuild) {
+          // Check for bins that are always-empty given a real profile.
+          if (hotness_bits == kHotBit) {
+            // This is not fatal, so only warn.
+            LOG(WARNING) << "Method " << method_ref.PrettyMethod() << " was hot but wasn't marked "
+                         << "either start-up or post-startup. Possible corrupted profile?";
+          }
+        }
+      }
 
       // Handle duplicate methods by pushing them repeatedly.
       OrderedMethodData method_data = {
-          method_hotness,
+          hotness_bits,
           oat_class,
           compiled_method,
           method_ref,
@@ -1279,6 +1278,10 @@
   }
 
  private:
+  // Cached profile index for the current dex file.
+  ProfileCompilationInfo::ProfileIndexType profile_index_;
+  const DexFile* profile_index_dex_file_;
+
   // List of compiled methods, later to be sorted by order defined in OrderedMethodData.
   // Methods can be inserted more than once in case of duplicated methods.
   OrderedMethodList ordered_methods_;
@@ -2358,7 +2361,7 @@
                   << "@ offset "
                   << relative_patcher_->GetOffset(ordered_method.method_reference)
                   << " X hotness "
-                  << reinterpret_cast<void*>(ordered_method.method_hotness.GetFlags());
+                  << ordered_method.hotness_bits;
       }
     }
   }
diff --git a/libprofile/profile/profile_compilation_info.cc b/libprofile/profile/profile_compilation_info.cc
index 8267174..5c95608 100644
--- a/libprofile/profile/profile_compilation_info.cc
+++ b/libprofile/profile/profile_compilation_info.cc
@@ -2263,21 +2263,6 @@
   return true;
 }
 
-template <typename Fn>
-ALWAYS_INLINE void ProfileCompilationInfo::DexFileData::ForMethodBitmapHotnessFlags(Fn fn) const {
-  uint32_t lastFlag = is_for_boot_image
-      ? MethodHotness::kFlagLastBoot
-      : MethodHotness::kFlagLastRegular;
-  for (uint32_t flag = MethodHotness::kFlagFirst; flag <= lastFlag; flag = flag << 1) {
-    if (flag == MethodHotness::kFlagHot) {
-      // There's no bit for hotness in the bitmap.
-      // We store the hotness by recording the method in the method list.
-      continue;
-    }
-    fn(enum_cast<MethodHotness::Flag>(flag));
-  }
-}
-
 void ProfileCompilationInfo::DexFileData::SetMethodHotness(size_t index,
                                                            MethodHotness::Flag flags) {
   DCHECK_LT(index, num_method_ids);
@@ -2286,6 +2271,7 @@
       method_bitmap.StoreBit(MethodFlagBitmapIndex(
           static_cast<MethodHotness::Flag>(flag), index), /*value=*/ true);
     }
+    return true;
   });
 }
 
@@ -2294,9 +2280,10 @@
   MethodHotness ret;
   ForMethodBitmapHotnessFlags([&](MethodHotness::Flag flag) {
     if (method_bitmap.LoadBit(MethodFlagBitmapIndex(
-          static_cast<MethodHotness::Flag>(flag), dex_method_index))) {
+            static_cast<MethodHotness::Flag>(flag), dex_method_index))) {
       ret.AddFlag(static_cast<MethodHotness::Flag>(flag));
     }
+    return true;
   });
   auto it = method_map.find(dex_method_index);
   if (it != method_map.end()) {
@@ -2325,24 +2312,6 @@
 static_assert(ProfileCompilationInfo::MethodHotness::kFlagStartupMaxBin == 1 << 15);
 static_assert(ProfileCompilationInfo::MethodHotness::kFlagLastBoot == 1 << 15);
 
-size_t ProfileCompilationInfo::DexFileData::MethodFlagBitmapIndex(
-      MethodHotness::Flag flag, size_t method_index) const {
-  DCHECK_LT(method_index, num_method_ids);
-  // The format is [startup bitmap][post startup bitmap][AmStartup][...]
-  // This compresses better than ([startup bit][post startup bit])*
-  return method_index + FlagBitmapIndex(flag) * num_method_ids;
-}
-
-size_t ProfileCompilationInfo::DexFileData::FlagBitmapIndex(MethodHotness::Flag flag) {
-  DCHECK(flag != MethodHotness::kFlagHot);
-  DCHECK(IsPowerOfTwo(static_cast<uint32_t>(flag)));
-  // We arrange the method flags in order, starting with the startup flag.
-  // The kFlagHot is not encoded in the bitmap and thus not expected as an
-  // argument here. Since all the other flags start at 1 we have to subtract
-  // one for the power of 2.
-  return WhichPowerOf2(static_cast<uint32_t>(flag)) - 1;
-}
-
 uint16_t ProfileCompilationInfo::DexFileData::GetUsedBitmapFlags() const {
   uint32_t used_flags = 0u;
   ForMethodBitmapHotnessFlags([&](MethodHotness::Flag flag) {
@@ -2350,6 +2319,7 @@
     if (method_bitmap.HasSomeBitSet(index * num_method_ids, num_method_ids)) {
       used_flags |= flag;
     }
+    return true;
   });
   return dchecked_integral_cast<uint16_t>(used_flags);
 }
@@ -2620,6 +2590,7 @@
       saved_bitmap.StoreBits(saved_bitmap_index * num_method_ids, src, num_method_ids);
       ++saved_bitmap_index;
     }
+    return true;
   });
   DCHECK_EQ(saved_bitmap_index * num_method_ids, saved_bitmap_bit_size);
   buffer.Advance(saved_bitmap_byte_size);
@@ -2723,6 +2694,7 @@
       method_bitmap.OrBits(index * num_method_ids, src, num_method_ids);
       ++saved_bitmap_index;
     }
+    return true;
   });
   buffer.Advance(saved_bitmap_byte_size);
 
diff --git a/libprofile/profile/profile_compilation_info.h b/libprofile/profile/profile_compilation_info.h
index 2061867..fe978d5 100644
--- a/libprofile/profile/profile_compilation_info.h
+++ b/libprofile/profile/profile_compilation_info.h
@@ -300,6 +300,15 @@
     return std::numeric_limits<ProfileIndexType>::max();
   }
 
+  // Find a tracked dex file. Returns `MaxProfileIndex()` on failure, whether due to no records
+  // for the dex location or profile key, or checksum/num_type_ids/num_method_ids mismatch.
+  ProfileIndexType FindDexFile(
+      const DexFile& dex_file,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const {
+    const DexFileData* data = FindDexDataUsingAnnotations(&dex_file, annotation);
+    return (data != nullptr) ? data->profile_index : MaxProfileIndex();
+  }
+
   // Find or add a tracked dex file. Returns `MaxProfileIndex()` on failure, whether due to
   // checksum/num_type_ids/num_method_ids mismatch or reaching the maximum number of dex files.
   ProfileIndexType FindOrAddDexFile(
@@ -485,6 +494,28 @@
   // Return the number of resolved classes that were profiled.
   uint32_t GetNumberOfResolvedClasses() const;
 
+  // Returns whether the referenced method is a startup method.
+  bool IsStartupMethod(ProfileIndexType profile_index, uint32_t method_index) const {
+    return info_[profile_index]->IsStartupMethod(method_index);
+  }
+
+  // Returns whether the referenced method is a post-startup method.
+  bool IsPostStartupMethod(ProfileIndexType profile_index, uint32_t method_index) const {
+    return info_[profile_index]->IsPostStartupMethod(method_index);
+  }
+
+  // Returns whether the referenced method is hot.
+  bool IsHotMethod(ProfileIndexType profile_index, uint32_t method_index) const {
+    return info_[profile_index]->IsHotMethod(method_index);
+  }
+
+  // Returns whether the referenced method is in the profile (with any hotness flag).
+  bool IsMethodInProfile(ProfileIndexType profile_index, uint32_t method_index) const {
+    DCHECK_LT(profile_index, info_.size());
+    const DexFileData* const data = info_[profile_index].get();
+    return data->IsMethodInProfile(method_index);
+  }
+
   // Returns the profile method info for a given method reference.
   //
   // Note that if the profile was built with annotations, the same dex file may be
@@ -499,6 +530,16 @@
       const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
 
   // Return true if the class's type is present in the profiling info.
+  bool ContainsClass(ProfileIndexType profile_index, dex::TypeIndex type_index) const {
+    DCHECK_LT(profile_index, info_.size());
+    const DexFileData* const data = info_[profile_index].get();
+    DCHECK(type_index.IsValid());
+    DCHECK(type_index.index_ <= data->num_type_ids ||
+           type_index.index_ - data->num_type_ids < extra_descriptors_.size());
+    return data->class_set.find(type_index) != data->class_set.end();
+  }
+
+  // Return true if the class's type is present in the profiling info.
   //
   // Note: see GetMethodHotness docs for the handling of annotations.
   bool ContainsClass(
@@ -782,6 +823,37 @@
     void SetMethodHotness(size_t index, MethodHotness::Flag flags);
     MethodHotness GetHotnessInfo(uint32_t dex_method_index) const;
 
+    bool IsStartupMethod(uint32_t method_index) const {
+      DCHECK_LT(method_index, num_method_ids);
+      return method_bitmap.LoadBit(
+          MethodFlagBitmapIndex(MethodHotness::Flag::kFlagStartup, method_index));
+    }
+
+    bool IsPostStartupMethod(uint32_t method_index) const {
+      DCHECK_LT(method_index, num_method_ids);
+      return method_bitmap.LoadBit(
+          MethodFlagBitmapIndex(MethodHotness::Flag::kFlagPostStartup, method_index));
+    }
+
+    bool IsHotMethod(uint32_t method_index) const {
+      DCHECK_LT(method_index, num_method_ids);
+      return method_map.find(method_index) != method_map.end();
+    }
+
+    bool IsMethodInProfile(uint32_t method_index) const {
+      DCHECK_LT(method_index, num_method_ids);
+      bool has_flag = false;
+      ForMethodBitmapHotnessFlags([&](MethodHotness::Flag flag) {
+        if (method_bitmap.LoadBit(MethodFlagBitmapIndex(
+                static_cast<MethodHotness::Flag>(flag), method_index))) {
+          has_flag = true;
+          return false;
+        }
+        return true;
+      });
+      return has_flag || IsHotMethod(method_index);
+    }
+
     bool ContainsClass(dex::TypeIndex type_index) const;
 
     uint32_t ClassesDataSize() const;
@@ -827,11 +899,41 @@
 
    private:
     template <typename Fn>
-    void ForMethodBitmapHotnessFlags(Fn fn) const;
+    void ForMethodBitmapHotnessFlags(Fn fn) const {
+      uint32_t lastFlag = is_for_boot_image
+          ? MethodHotness::kFlagLastBoot
+          : MethodHotness::kFlagLastRegular;
+      for (uint32_t flag = MethodHotness::kFlagFirst; flag <= lastFlag; flag = flag << 1) {
+        if (flag == MethodHotness::kFlagHot) {
+          // There's no bit for hotness in the bitmap.
+          // We store the hotness by recording the method in the method list.
+          continue;
+        }
+        bool cont = fn(enum_cast<MethodHotness::Flag>(flag));
+        if (!cont) {
+          break;
+        }
+      }
+    }
+
+    size_t MethodFlagBitmapIndex(MethodHotness::Flag flag, size_t method_index) const {
+      DCHECK_LT(method_index, num_method_ids);
+      // The format is [startup bitmap][post startup bitmap][AmStartup][...]
+      // This compresses better than ([startup bit][post startup bit])*
+      return method_index + FlagBitmapIndex(flag) * num_method_ids;
+    }
+
+    static size_t FlagBitmapIndex(MethodHotness::Flag flag) {
+      DCHECK(flag != MethodHotness::kFlagHot);
+      DCHECK(IsPowerOfTwo(static_cast<uint32_t>(flag)));
+      // We arrange the method flags in order, starting with the startup flag.
+      // The kFlagHot is not encoded in the bitmap and thus not expected as an
+      // argument here. Since all the other flags start at 1 we have to subtract
+      // one from the power of 2.
+      return WhichPowerOf2(static_cast<uint32_t>(flag)) - 1;
+    }
 
     static void WriteClassSet(SafeBuffer& buffer, const ArenaSet<dex::TypeIndex>& class_set);
-    size_t MethodFlagBitmapIndex(MethodHotness::Flag flag, size_t method_index) const;
-    static size_t FlagBitmapIndex(MethodHotness::Flag flag);
 
     uint16_t GetUsedBitmapFlags() const;
   };