Introduce the ability to annotate profile samples

We can now annotate profiles samples (classes or methods) with additional
metadata that will be persisted in the offline representation. Currently
the annotations support the package name that contributed the given
samples.

When samples are annotated, they are grouped into distinct categories
indexed by (dex_file, sample_annotation). This is achieved by extending
the profile key to include a serialized representation of the annotation.

Because they create independent groups in the profile, the annotations can
potentially increase the profile size considerably so care should be taken
when adding them in big numbers.

Information extraction (methods and classes) has also been extended to
support the annotations. Users may choose to extract the info for
a particular group (dex_file, sample_annotation) or, as before, just for a
single dex file. If the metadata is not given (e.g. when using profile
guided compilation), the default search mechanism kicks in, and the first
dex file matching the constraint is searched.

By extending the key representation, we preserve the previous profile
behaviour without the need to extended the underlying format or increase
the version.

Bug: 139884006
Test: m test-art-host-gtest

Change-Id: Iaccecd05c575bf0dac6dace6257cdafc6dc4a329
diff --git a/libprofile/profile/profile_compilation_info.cc b/libprofile/profile/profile_compilation_info.cc
index 013d5cf..aa44b64 100644
--- a/libprofile/profile/profile_compilation_info.cc
+++ b/libprofile/profile/profile_compilation_info.cc
@@ -68,6 +68,15 @@
 // DO NOT CHANGE THIS! (it's similar to classes.dex in the apk files).
 const char ProfileCompilationInfo::kDexMetadataProfileEntry[] = "primary.prof";
 
+// A synthetic annotations that can be used to denote that no annotation should
+// be associated with the profile samples. We use the empty string for the package name
+// because that's an invalid package name and should never occur in practice.
+const ProfileCompilationInfo::ProfileSampleAnnotation
+  ProfileCompilationInfo::ProfileSampleAnnotation::kNone =
+      ProfileCompilationInfo::ProfileSampleAnnotation("");
+
+static constexpr char kSampleMetdataSeparator = ':';
+
 static constexpr uint16_t kMaxDexFileKeyLength = PATH_MAX;
 
 // Debug flag to ignore checksums when testing if a method or a class is present in the profile.
@@ -149,12 +158,13 @@
 
 // Transform the actual dex location into a key used to index the dex file in the profile.
 // See ProfileCompilationInfo#GetProfileDexFileBaseKey as well.
-// For regular profiles (non-boot) the profile key is the same as its base key.
-std::string ProfileCompilationInfo::GetProfileDexFileKey(const std::string& dex_location) const {
-  // TODO(calin): append the RuntimeInstructionSet to the key so we can capture arch-dependent data.
-  // This requires a bit of work as the ProfileSaver and various tests rely on this being a
-  // static public method.
-  return GetProfileDexFileBaseKey(dex_location);
+std::string ProfileCompilationInfo::GetProfileDexFileAugmentedKey(
+      const std::string& dex_location,
+      const ProfileSampleAnnotation& annotation) {
+  std::string base_key = GetProfileDexFileBaseKey(dex_location);
+  return annotation == ProfileSampleAnnotation::kNone
+      ? base_key
+      : base_key + kSampleMetdataSeparator + annotation.GetOriginPackageName();;
 }
 
 // Transform the actual dex location into a base profile key (represented as relative paths).
@@ -172,10 +182,26 @@
   }
 }
 
+std::string ProfileCompilationInfo::GetBaseKeyFromAugmentedKey(
+    const std::string& profile_key) {
+  size_t pos = profile_key.rfind(kSampleMetdataSeparator);
+  return (pos == std::string::npos) ? profile_key : profile_key.substr(0, pos);
+}
+
+std::string ProfileCompilationInfo::MigrateAnnotationInfo(
+    const std::string& base_key,
+    const std::string& augmented_key) {
+  size_t pos = augmented_key.rfind(kSampleMetdataSeparator);
+  return (pos == std::string::npos)
+      ? base_key
+      : base_key + augmented_key.substr(pos);
+}
+
 bool ProfileCompilationInfo::AddMethods(const std::vector<ProfileMethodInfo>& methods,
-                                        MethodHotness::Flag flags) {
+                                        MethodHotness::Flag flags,
+                                        const ProfileSampleAnnotation& annotation) {
   for (const ProfileMethodInfo& method : methods) {
-    if (!AddMethod(method, flags)) {
+    if (!AddMethod(method, flags, annotation)) {
       return false;
     }
   }
@@ -641,8 +667,31 @@
   return result;
 }
 
-bool ProfileCompilationInfo::AddMethod(const ProfileMethodInfo& pmi, MethodHotness::Flag flags) {
-  DexFileData* const data = GetOrAddDexFileData(pmi.ref.dex_file);
+const ProfileCompilationInfo::DexFileData* ProfileCompilationInfo::FindDexDataUsingAnnotations(
+      const DexFile* dex_file,
+      const ProfileSampleAnnotation& annotation) const {
+  if (annotation == ProfileSampleAnnotation::kNone) {
+    std::string profile_key = GetProfileDexFileBaseKey(dex_file->GetLocation());
+    for (const DexFileData* dex_data : info_) {
+      if (profile_key == GetBaseKeyFromAugmentedKey(dex_data->profile_key)) {
+        if (!ChecksumMatch(dex_data->checksum, dex_file->GetLocationChecksum())) {
+          return nullptr;
+        }
+        return dex_data;
+      }
+    }
+  } else {
+    std::string profile_key = GetProfileDexFileAugmentedKey(dex_file->GetLocation(), annotation);
+    return FindDexData(profile_key, dex_file->GetLocationChecksum());
+  }
+
+  return nullptr;
+}
+
+bool ProfileCompilationInfo::AddMethod(const ProfileMethodInfo& pmi,
+                                       MethodHotness::Flag flags,
+                                       const ProfileSampleAnnotation& annotation) {
+  DexFileData* const data = GetOrAddDexFileData(pmi.ref.dex_file, annotation);
   if (data == nullptr) {  // checksum mismatch
     return false;
   }
@@ -664,7 +713,7 @@
       continue;
     }
     for (const TypeReference& class_ref : cache.classes) {
-      DexFileData* class_dex_data = GetOrAddDexFileData(class_ref.dex_file);
+      DexFileData* class_dex_data = GetOrAddDexFileData(class_ref.dex_file, annotation);
       if (class_dex_data == nullptr) {  // checksum mismatch
         return false;
       }
@@ -1021,10 +1070,11 @@
 bool ProfileCompilationInfo::VerifyProfileData(const std::vector<const DexFile*>& dex_files) {
   std::unordered_map<std::string, const DexFile*> key_to_dex_file;
   for (const DexFile* dex_file : dex_files) {
-    key_to_dex_file.emplace(GetProfileDexFileKey(dex_file->GetLocation()), dex_file);
+    key_to_dex_file.emplace(GetProfileDexFileBaseKey(dex_file->GetLocation()), dex_file);
   }
   for (const DexFileData* dex_data : info_) {
-    const auto it = key_to_dex_file.find(dex_data->profile_key);
+    // We need to remove any annotation from the key during verification.
+    const auto it = key_to_dex_file.find(GetBaseKeyFromAugmentedKey(dex_data->profile_key));
     if (it == key_to_dex_file.end()) {
       // It is okay if profile contains data for additional dex files.
       continue;
@@ -1506,23 +1556,19 @@
   return true;
 }
 
-const ProfileCompilationInfo::DexFileData* ProfileCompilationInfo::FindDexData(
-    const DexFile* dex_file) const {
-  return FindDexData(GetProfileDexFileKey(dex_file->GetLocation()),
-                     dex_file->GetLocationChecksum());
-}
-
 ProfileCompilationInfo::MethodHotness ProfileCompilationInfo::GetMethodHotness(
-    const MethodReference& method_ref) const {
-  const DexFileData* dex_data = FindDexData(method_ref.dex_file);
+    const MethodReference& method_ref,
+    const ProfileSampleAnnotation& annotation) const {
+  const DexFileData* dex_data = FindDexDataUsingAnnotations(method_ref.dex_file, annotation);
   return dex_data != nullptr
       ? dex_data->GetHotnessInfo(method_ref.index)
       : MethodHotness();
 }
 
 std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo>
-ProfileCompilationInfo::GetHotMethodInfo(const MethodReference& method_ref) const {
-  MethodHotness hotness(GetMethodHotness(method_ref));
+ProfileCompilationInfo::GetHotMethodInfo(const MethodReference& method_ref,
+                                         const ProfileSampleAnnotation& annotation) const {
+  MethodHotness hotness(GetMethodHotness(method_ref, annotation));
   if (!hotness.IsHot()) {
     return nullptr;
   }
@@ -1541,8 +1587,10 @@
 }
 
 
-bool ProfileCompilationInfo::ContainsClass(const DexFile& dex_file, dex::TypeIndex type_idx) const {
-  const DexFileData* dex_data = FindDexData(&dex_file);
+bool ProfileCompilationInfo::ContainsClass(const DexFile& dex_file,
+                                           dex::TypeIndex type_idx,
+                                           const ProfileSampleAnnotation& annotation) const {
+  const DexFileData* dex_data = FindDexDataUsingAnnotations(&dex_file, annotation);
   return (dex_data != nullptr) && dex_data->ContainsClass(type_idx);
 }
 
@@ -1579,14 +1627,15 @@
       os << dex_data->profile_key;
     } else {
       // Replace the (empty) multidex suffix of the first key with a substitute for easier reading.
-      std::string multidex_suffix = DexFileLoader::GetMultiDexSuffix(dex_data->profile_key);
+      std::string multidex_suffix = DexFileLoader::GetMultiDexSuffix(
+          GetBaseKeyFromAugmentedKey(dex_data->profile_key));
       os << (multidex_suffix.empty() ? kFirstDexFileKeySubstitute : multidex_suffix);
     }
     os << " [index=" << static_cast<uint32_t>(dex_data->profile_index) << "]";
     os << " [checksum=" << std::hex << dex_data->checksum << "]" << std::dec;
     const DexFile* dex_file = nullptr;
     for (const DexFile* current : dex_files) {
-      if (dex_data->profile_key == current->GetLocation() &&
+      if (GetBaseKeyFromAugmentedKey(dex_data->profile_key) == current->GetLocation() &&
           dex_data->checksum == current->GetLocationChecksum()) {
         dex_file = current;
       }
@@ -1651,9 +1700,10 @@
     /*out*/std::set<dex::TypeIndex>* class_set,
     /*out*/std::set<uint16_t>* hot_method_set,
     /*out*/std::set<uint16_t>* startup_method_set,
-    /*out*/std::set<uint16_t>* post_startup_method_method_set) const {
+    /*out*/std::set<uint16_t>* post_startup_method_method_set,
+    const ProfileSampleAnnotation& annotation) const {
   std::set<std::string> ret;
-  const DexFileData* dex_data = FindDexData(&dex_file);
+  const DexFileData* dex_data = FindDexDataUsingAnnotations(&dex_file, annotation);
   if (dex_data == nullptr) {
     return false;
   }
@@ -1722,7 +1772,7 @@
 
   for (uint16_t i = 0; i < number_of_dex_files; i++) {
     std::string dex_location = DexFileLoader::GetMultiDexLocation(i, base_dex_location.c_str());
-    std::string profile_key = info.GetProfileDexFileKey(dex_location);
+    std::string profile_key = info.GetProfileDexFileBaseKey(dex_location);
 
     DexFileData* const data = info.GetOrAddDexFileData(profile_key, /*checksum=*/ 0, max_method);
     for (uint16_t m = 0; m < number_of_methods; m++) {
@@ -2021,10 +2071,11 @@
 }
 
 HashSet<std::string> ProfileCompilationInfo::GetClassDescriptors(
-    const std::vector<const DexFile*>& dex_files) {
+    const std::vector<const DexFile*>& dex_files,
+    const ProfileSampleAnnotation& annotation) {
   HashSet<std::string> ret;
   for (const DexFile* dex_file : dex_files) {
-    const DexFileData* data = FindDexData(dex_file);
+    const DexFileData* data = FindDexDataUsingAnnotations(dex_file, annotation);
     if (data != nullptr) {
       for (dex::TypeIndex type_idx : data->class_set) {
         if (!dex_file->IsTypeIndexValid(type_idx)) {
@@ -2079,8 +2130,9 @@
     for (DexFileData* dex_data : info_) {
       if (dex_data->checksum == dex_file->GetLocationChecksum()
           && dex_data->num_method_ids == dex_file->NumMethodIds()) {
-        std::string new_profile_key = GetProfileDexFileKey(dex_file->GetLocation());
-        if (dex_data->profile_key != new_profile_key) {
+        std::string new_profile_key = GetProfileDexFileBaseKey(dex_file->GetLocation());
+        std::string dex_data_base_key = GetBaseKeyFromAugmentedKey(dex_data->profile_key);
+        if (dex_data_base_key != new_profile_key) {
           if (profile_key_map_.find(new_profile_key) != profile_key_map_.end()) {
             // We can't update the key if the new key belongs to a different dex file.
             LOG(ERROR) << "Cannot update profile key to " << new_profile_key
@@ -2088,7 +2140,10 @@
             return false;
           }
           profile_key_map_.erase(dex_data->profile_key);
-          profile_key_map_.Put(new_profile_key, dex_data->profile_index);
+          // 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;
         }
       }
@@ -2140,4 +2195,8 @@
   return stream;
 }
 
+bool ProfileCompilationInfo::ProfileSampleAnnotation::operator==(
+      const ProfileSampleAnnotation& other) const {
+  return origin_package_name_ == other.origin_package_name_;
+}
 }  // namespace art
diff --git a/libprofile/profile/profile_compilation_info.h b/libprofile/profile/profile_compilation_info.h
index 55be2b6..7043584 100644
--- a/libprofile/profile/profile_compilation_info.h
+++ b/libprofile/profile/profile_compilation_info.h
@@ -271,13 +271,37 @@
         : inline_caches(inline_cache_map) {}
 
     bool operator==(const OfflineProfileMethodInfo& other) const;
-    // Checks that this offline represenation of inline caches matches the runtime view of the data.
+    // Checks that this offline representation of inline caches matches the runtime view of the
+    // data.
     bool operator==(const std::vector<ProfileMethodInfo::ProfileInlineCache>& other) const;
 
     const InlineCacheMap* const inline_caches;
     std::vector<DexReference> dex_references;
   };
 
+  // Encapsulates metadata that can be associated with the methods and classes added to the profile.
+  // The additional metadata is serialized in the profile and becomes part of the profile key
+  // representation. It can be used to differentiate the samples that are added to the profile
+  // based on the supported criteria (e.g. keep track of which app generated what sample when
+  // constructing a boot profile.).
+  class ProfileSampleAnnotation {
+   public:
+    explicit ProfileSampleAnnotation(const std::string& package_name) :
+        origin_package_name_(package_name) {}
+
+    const std::string& GetOriginPackageName() const { return origin_package_name_; }
+
+    bool operator==(const ProfileSampleAnnotation& other) const;
+
+    // A convenient empty annotation object that can be used to denote that no annotation should
+    // be associated with the profile samples.
+    static const ProfileSampleAnnotation kNone;
+
+   private:
+    // The name of the package that generated the samples.
+    const std::string origin_package_name_;
+  };
+
   // Public methods to create, extend or query the profile.
   ProfileCompilationInfo();
   explicit ProfileCompilationInfo(bool for_boot_image);
@@ -287,13 +311,24 @@
   ~ProfileCompilationInfo();
 
   // Add the given methods to the current profile object.
-  bool AddMethods(const std::vector<ProfileMethodInfo>& methods, MethodHotness::Flag flags);
+  //
+  // Note: if an annotation is provided, the methods/classes will be associated with the group
+  // (dex_file, sample_annotation). Each group keeps its unique set of methods/classes.
+  bool AddMethods(const std::vector<ProfileMethodInfo>& methods,
+                  MethodHotness::Flag flags,
+                  const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone);
 
   // Add multiple type ids for classes in a single dex file. Iterator is for type_ids not
   // class_defs.
+  //
+  // Note: see AddMethods docs for the handling of annotations.
   template <class Iterator>
-  bool AddClassesForDex(const DexFile* dex_file, Iterator index_begin, Iterator index_end) {
-    DexFileData* data = GetOrAddDexFileData(dex_file);
+  bool AddClassesForDex(
+      const DexFile* dex_file,
+      Iterator index_begin,
+      Iterator index_end,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) {
+    DexFileData* data = GetOrAddDexFileData(dex_file, annotation);
     if (data == nullptr) {
       return false;
     }
@@ -302,16 +337,24 @@
   }
 
   // Add a method to the profile using its online representation (containing runtime structures).
-  bool AddMethod(const ProfileMethodInfo& pmi, MethodHotness::Flag flags);
+  //
+  // Note: see AddMethods docs for the handling of annotations.
+  bool AddMethod(const ProfileMethodInfo& pmi,
+                 MethodHotness::Flag flags,
+                 const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone);
 
   // Bulk add sampled methods and/or hot methods for a single dex, fast since it only has one
   // GetOrAddDexFileData call.
+  //
+  // Note: see AddMethods docs for the handling of annotations.
   template <class Iterator>
-  bool AddMethodsForDex(MethodHotness::Flag flags,
-                        const DexFile* dex_file,
-                        Iterator index_begin,
-                        Iterator index_end) {
-    DexFileData* data = GetOrAddDexFileData(dex_file);
+  bool AddMethodsForDex(
+      MethodHotness::Flag flags,
+      const DexFile* dex_file,
+      Iterator index_begin,
+      Iterator index_end,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) {
+    DexFileData* data = GetOrAddDexFileData(dex_file, annotation);
     if (data == nullptr) {
       return false;
     }
@@ -377,17 +420,35 @@
   uint32_t GetNumberOfResolvedClasses() const;
 
   // Returns the profile method info for a given method reference.
-  MethodHotness GetMethodHotness(const MethodReference& method_ref) const;;
+  //
+  // Note that if the profile was built with annotations, the same dex file may be
+  // represented multiple times in the profile (due to different annotation associated with it).
+  // If so, and if no annotation is passed to this method, then only the first dex file is searched.
+  //
+  // Implementation details: It is suitable to pass kNone for regular profile guided compilation
+  // because during compilation we generally don't care about annotations. The metadata is
+  // useful for boot profiles which need the extra information.
+  MethodHotness GetMethodHotness(
+      const MethodReference& method_ref,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
 
   // Return true if the class's type is present in the profiling info.
-  bool ContainsClass(const DexFile& dex_file, dex::TypeIndex type_idx) const;
+  //
+  // Note: see GetMethodHotness docs for the handling of annotations.
+  bool ContainsClass(
+      const DexFile& dex_file,
+      dex::TypeIndex type_idx,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
 
   // Return the hot method info for the given location and index from the profiling info.
   // If the method index is not found or the checksum doesn't match, null is returned.
   // Note: the inline cache map is a pointer to the map stored in the profile and
   // its allocation will go away if the profile goes out of scope.
+  //
+  // Note: see GetMethodHotness docs for the handling of annotations.
   std::unique_ptr<OfflineProfileMethodInfo> GetHotMethodInfo(
-      const MethodReference& method_ref) const;
+      const MethodReference& method_ref,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
 
   // Dump all the loaded profile info into a string and returns it.
   // If dex_files is not empty then the method indices will be resolved to their
@@ -399,11 +460,15 @@
   // Return the classes and methods for a given dex file through out args. The out args are the set
   // of class as well as the methods and their associated inline caches. Returns true if the dex
   // file is register and has a matching checksum, false otherwise.
-  bool GetClassesAndMethods(const DexFile& dex_file,
-                            /*out*/std::set<dex::TypeIndex>* class_set,
-                            /*out*/std::set<uint16_t>* hot_method_set,
-                            /*out*/std::set<uint16_t>* startup_method_set,
-                            /*out*/std::set<uint16_t>* post_startup_method_method_set) const;
+  //
+  // Note: see GetMethodHotness docs for the handling of annotations.
+  bool GetClassesAndMethods(
+      const DexFile& dex_file,
+      /*out*/std::set<dex::TypeIndex>* class_set,
+      /*out*/std::set<uint16_t>* hot_method_set,
+      /*out*/std::set<uint16_t>* startup_method_set,
+      /*out*/std::set<uint16_t>* post_startup_method_method_set,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
 
   // Returns true iff both profiles have the same version.
   bool SameVersion(const ProfileCompilationInfo& other) const;
@@ -411,11 +476,9 @@
   // Perform an equality test with the `other` profile information.
   bool Equals(const ProfileCompilationInfo& other);
 
-  // Return the profile key associated with the given dex location.
-  std::string GetProfileDexFileKey(const std::string& dex_location) const;
   // Return the base profile key associated with the given dex location. The base profile key
   // is solely constructed based on the dex location (as opposed to the one produced by
-  // GetProfileDexFileKey which may include additional metadata like architecture or source
+  // GetProfileDexFileAugmentedKey which may include additional metadata like the origin
   // package name)
   static std::string GetProfileDexFileBaseKey(const std::string& dex_location);
 
@@ -442,7 +505,10 @@
   ArenaAllocator* GetAllocator() { return &allocator_; }
 
   // Return all of the class descriptors in the profile for a set of dex files.
-  HashSet<std::string> GetClassDescriptors(const std::vector<const DexFile*>& dex_files);
+  // Note: see GetMethodHotness docs for the handling of annotations..
+  HashSet<std::string> GetClassDescriptors(
+      const std::vector<const DexFile*>& dex_files,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone);
 
   // Return true if the fd points to a profile file.
   bool IsProfileFile(int fd);
@@ -585,8 +651,9 @@
                                    uint32_t checksum,
                                    uint32_t num_method_ids);
 
-  DexFileData* GetOrAddDexFileData(const DexFile* dex_file) {
-    return GetOrAddDexFileData(GetProfileDexFileKey(dex_file->GetLocation()),
+  DexFileData* GetOrAddDexFileData(const DexFile* dex_file,
+                                   const ProfileSampleAnnotation& annotation) {
+    return GetOrAddDexFileData(GetProfileDexFileAugmentedKey(dex_file->GetLocation(), annotation),
                                dex_file->GetLocationChecksum(),
                                dex_file->NumMethodIds());
   }
@@ -600,10 +667,14 @@
   const DexFileData* FindDexData(const std::string& profile_key,
                                  uint32_t checksum,
                                  bool verify_checksum = true) const;
-
-  // Return the dex data associated with the given dex file or null if the profile doesn't contain
-  // the key or the checksum mismatches.
-  const DexFileData* FindDexData(const DexFile* dex_file) const;
+  // Same as FindDexData but performs the searching using the given annotation:
+  //   - If the annotation is kNone then the search ignores it and only looks at the base keys.
+  //     In this case only the first matching dex is searched.
+  //   - If the annotation is not kNone, the augmented key is constructed and used to invoke
+  //     the regular FindDexData.
+  const DexFileData* FindDexDataUsingAnnotations(
+      const DexFile* dex_file,
+      const ProfileSampleAnnotation& annotation) const;
 
   // Inflate the input buffer (in_buffer) of size in_size. It returns a buffer of
   // compressed data for the input buffer of "compressed_data_size" size.
@@ -811,6 +882,21 @@
   // Returns the threshold size (in bytes) which will cause save/load failures.
   size_t GetSizeErrorThresholdBytes() const;
 
+
+  // Returns the augmented profile key associated with the given dex location.
+  // The return key will contain a serialized form of the information from the provided
+  // annotation. If the annotation is ProfileSampleAnnotation::kNone then no extra info is
+  // added to the key and this method is equivalent to GetProfileDexFileBaseKey.
+  static std::string GetProfileDexFileAugmentedKey(const std::string& dex_location,
+                                                   const ProfileSampleAnnotation& annotation);
+
+  // Returns a base key without the annotation information.
+  static std::string GetBaseKeyFromAugmentedKey(const std::string& profile_key);
+
+  // Migrates the annotation from an augmented key to a base key.
+  static std::string MigrateAnnotationInfo(const std::string& base_key,
+                                           const std::string& augmented_key);
+
   friend class ProfileCompilationInfoTest;
   friend class CompilerDriverProfileTest;
   friend class ProfileAssistantTest;
diff --git a/libprofile/profile/profile_compilation_info_test.cc b/libprofile/profile/profile_compilation_info_test.cc
index 7b4ad5a..6390b39 100644
--- a/libprofile/profile/profile_compilation_info_test.cc
+++ b/libprofile/profile/profile_compilation_info_test.cc
@@ -33,6 +33,7 @@
 
 using Hotness = ProfileCompilationInfo::MethodHotness;
 using ProfileInlineCache = ProfileMethodInfo::ProfileInlineCache;
+using ProfileSampleAnnotation = ProfileCompilationInfo::ProfileSampleAnnotation;
 
 static constexpr size_t kMaxMethodIds = 65535;
 static uint32_t kMaxHotnessFlagBootIndex =
@@ -67,32 +68,31 @@
  protected:
   bool AddMethod(ProfileCompilationInfo* info,
                  const DexFile* dex,
-                 uint16_t method_idx) {
-    return AddMethod(info, dex, method_idx, Hotness::kFlagHot);
-  }
-
-  bool AddMethod(ProfileCompilationInfo* info,
-                 const DexFile* dex,
                  uint16_t method_idx,
-                 Hotness::Flag flags) {
+                 Hotness::Flag flags = Hotness::kFlagHot,
+                 const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) {
     return info->AddMethod(ProfileMethodInfo(MethodReference(dex, method_idx)),
-                           flags);
+                           flags,
+                           annotation);
   }
 
   bool AddMethod(ProfileCompilationInfo* info,
                 const DexFile* dex,
                 uint16_t method_idx,
-                const std::vector<ProfileInlineCache>& inline_caches) {
+                const std::vector<ProfileInlineCache>& inline_caches,
+                const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) {
     return info->AddMethod(
         ProfileMethodInfo(MethodReference(dex, method_idx), inline_caches),
-        Hotness::kFlagHot);
+        Hotness::kFlagHot,
+        annotation);
   }
 
   bool AddClass(ProfileCompilationInfo* info,
                 const DexFile* dex,
-                dex::TypeIndex type_index) {
+                dex::TypeIndex type_index,
+                const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) {
     std::vector<dex::TypeIndex> classes = {type_index};
-    return info->AddClassesForDex(dex, classes.begin(), classes.end());
+    return info->AddClassesForDex(dex, classes.begin(), classes.end(), annotation);
   }
 
   uint32_t GetFd(const ScratchFile& file) {
@@ -102,8 +102,9 @@
   std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> GetMethod(
       const ProfileCompilationInfo& info,
       const DexFile* dex,
-      uint16_t method_idx) {
-    return info.GetHotMethodInfo(MethodReference(dex, method_idx));
+      uint16_t method_idx,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) {
+    return info.GetHotMethodInfo(MethodReference(dex, method_idx), annotation);
   }
 
   // Creates an inline cache which will be destructed at the end of the test.
@@ -1415,4 +1416,257 @@
   std::vector<ProfileMethodInfo> pmis = {ProfileMethodInfo(hot), ProfileMethodInfo(bad_ref)};
   ASSERT_FALSE(info.AddMethods(pmis, Hotness::kFlagHot));
 }
+
+// Verify that we can add methods with annotations.
+TEST_F(ProfileCompilationInfoTest, AddAnnotationsToMethods) {
+  ProfileCompilationInfo info;
+
+  ProfileSampleAnnotation psa1("test1");
+  ProfileSampleAnnotation psa2("test2");
+  // Save a few methods using different annotations, some overlapping, some not.
+  for (uint16_t i = 0; i < 10; i++) {
+    ASSERT_TRUE(AddMethod(&info, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa1));
+  }
+  for (uint16_t i = 5; i < 15; i++) {
+    ASSERT_TRUE(AddMethod(&info, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa2));
+  }
+
+  auto run_test = [&dex1 = dex1, &psa1 = psa1, &psa2 = psa2](const ProfileCompilationInfo& info) {
+    // Check that all methods are in.
+    for (uint16_t i = 0; i < 10; i++) {
+      EXPECT_TRUE(info.GetMethodHotness(MethodReference(dex1, i), psa1).IsInProfile());
+      EXPECT_TRUE(info.GetHotMethodInfo(MethodReference(dex1, i), psa1) != nullptr);
+    }
+    for (uint16_t i = 5; i < 15; i++) {
+      EXPECT_TRUE(info.GetMethodHotness(MethodReference(dex1, i), psa2).IsInProfile());
+      EXPECT_TRUE(info.GetHotMethodInfo(MethodReference(dex1, i), psa2) != nullptr);
+    }
+    // Check that the non-overlapping methods are not added with a wrong annotation.
+    for (uint16_t i = 10; i < 15; i++) {
+      EXPECT_FALSE(info.GetMethodHotness(MethodReference(dex1, i), psa1).IsInProfile());
+      EXPECT_FALSE(info.GetHotMethodInfo(MethodReference(dex1, i), psa1) != nullptr);
+    }
+    for (uint16_t i = 0; i < 5; i++) {
+      EXPECT_FALSE(info.GetMethodHotness(MethodReference(dex1, i), psa2).IsInProfile());
+      EXPECT_FALSE(info.GetHotMethodInfo(MethodReference(dex1, i), psa2) != nullptr);
+    }
+    // Check that when querying without an annotation only the first one is searched.
+    for (uint16_t i = 0; i < 10; i++) {
+      EXPECT_TRUE(info.GetMethodHotness(MethodReference(dex1, i)).IsInProfile());
+      EXPECT_TRUE(info.GetHotMethodInfo(MethodReference(dex1, i)) != nullptr);
+    }
+    // ... this should be false because they belong the second appearance of dex1.
+    for (uint16_t i = 10; i < 15; i++) {
+      EXPECT_FALSE(info.GetMethodHotness(MethodReference(dex1, i)).IsInProfile());
+      EXPECT_FALSE(info.GetHotMethodInfo(MethodReference(dex1, i)) != nullptr);
+    }
+
+    // Sanity check that methods cannot be found with a non existing annotation.
+    MethodReference ref(dex1, 0);
+    ProfileSampleAnnotation not_exisiting("A");
+    EXPECT_FALSE(info.GetMethodHotness(ref, not_exisiting).IsInProfile());
+    EXPECT_FALSE(info.GetHotMethodInfo(ref, not_exisiting) != nullptr);
+  };
+
+  // Run the test before save.
+  run_test(info);
+
+  ScratchFile profile;
+  ASSERT_TRUE(info.Save(GetFd(profile)));
+  ASSERT_EQ(0, profile.GetFile()->Flush());
+
+  // Check that we get back what we saved.
+  ProfileCompilationInfo loaded_info;
+  ASSERT_TRUE(profile.GetFile()->ResetOffset());
+  ASSERT_TRUE(loaded_info.Load(GetFd(profile)));
+  ASSERT_TRUE(loaded_info.Equals(info));
+
+  // Run the test after save and load.
+  run_test(loaded_info);
+}
+
+// Verify that we can add classes with annotations.
+TEST_F(ProfileCompilationInfoTest, AddAnnotationsToClasses) {
+  ProfileCompilationInfo info;
+
+  ProfileSampleAnnotation psa1("test1");
+  ProfileSampleAnnotation psa2("test2");
+  // Save a few classes using different annotations, some overlapping, some not.
+  for (uint16_t i = 0; i < 10; i++) {
+    ASSERT_TRUE(AddClass(&info, dex1, dex::TypeIndex(i), psa1));
+  }
+  for (uint16_t i = 5; i < 15; i++) {
+    ASSERT_TRUE(AddClass(&info, dex1, dex::TypeIndex(i), psa2));
+  }
+
+  auto run_test = [&dex1 = dex1, &psa1 = psa1, &psa2 = psa2](const ProfileCompilationInfo& info) {
+    // Check that all classes are in.
+    for (uint16_t i = 0; i < 10; i++) {
+      EXPECT_TRUE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa1));
+    }
+    for (uint16_t i = 5; i < 15; i++) {
+      EXPECT_TRUE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa2));
+    }
+    // Check that the non-overlapping classes are not added with a wrong annotation.
+    for (uint16_t i = 10; i < 15; i++) {
+      EXPECT_FALSE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa1));
+    }
+    for (uint16_t i = 0; i < 5; i++) {
+      EXPECT_FALSE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa2));
+    }
+    // Check that when querying without an annotation only the first one is searched.
+    for (uint16_t i = 0; i < 10; i++) {
+      EXPECT_TRUE(info.ContainsClass(*dex1, dex::TypeIndex(i)));
+    }
+    // ... this should be false because they belong the second appearance of dex1.
+    for (uint16_t i = 10; i < 15; i++) {
+      EXPECT_FALSE(info.ContainsClass(*dex1, dex::TypeIndex(i)));
+    }
+
+    // Sanity check that classes cannot be found with a non existing annotation.
+    EXPECT_FALSE(info.ContainsClass(*dex1, dex::TypeIndex(0), ProfileSampleAnnotation("new_test")));
+  };
+
+  // Run the test before save.
+  run_test(info);
+
+  ScratchFile profile;
+  ASSERT_TRUE(info.Save(GetFd(profile)));
+  ASSERT_EQ(0, profile.GetFile()->Flush());
+
+  // Check that we get back what we saved.
+  ProfileCompilationInfo loaded_info;
+  ASSERT_TRUE(profile.GetFile()->ResetOffset());
+  ASSERT_TRUE(loaded_info.Load(GetFd(profile)));
+  ASSERT_TRUE(loaded_info.Equals(info));
+
+  // Run the test after save and load.
+  run_test(loaded_info);
+}
+
+// Verify we can merge samples with annotations.
+TEST_F(ProfileCompilationInfoTest, MergeWithAnnotations) {
+  ProfileCompilationInfo info1;
+  ProfileCompilationInfo info2;
+
+  ProfileSampleAnnotation psa1("test1");
+  ProfileSampleAnnotation psa2("test2");
+
+  for (uint16_t i = 0; i < 10; i++) {
+    ASSERT_TRUE(AddMethod(&info1, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa1));
+    ASSERT_TRUE(AddClass(&info1, dex1, dex::TypeIndex(i), psa1));
+  }
+  for (uint16_t i = 5; i < 15; i++) {
+    ASSERT_TRUE(AddMethod(&info2, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa1));
+    ASSERT_TRUE(AddMethod(&info2, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa2));
+    ASSERT_TRUE(AddMethod(&info2, dex2, /* method_idx= */ i, Hotness::kFlagHot, psa2));
+    ASSERT_TRUE(AddClass(&info2, dex1, dex::TypeIndex(i), psa1));
+    ASSERT_TRUE(AddClass(&info2, dex1, dex::TypeIndex(i), psa2));
+  }
+
+  ProfileCompilationInfo info;
+  ASSERT_TRUE(info.MergeWith(info1));
+  ASSERT_TRUE(info.MergeWith(info2));
+
+  // Check that all items are in.
+  for (uint16_t i = 0; i < 15; i++) {
+    EXPECT_TRUE(info.GetMethodHotness(MethodReference(dex1, i), psa1).IsInProfile());
+    EXPECT_TRUE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa1));
+  }
+  for (uint16_t i = 5; i < 15; i++) {
+    EXPECT_TRUE(info.GetMethodHotness(MethodReference(dex1, i), psa2).IsInProfile());
+    EXPECT_TRUE(info.GetMethodHotness(MethodReference(dex2, i), psa2).IsInProfile());
+    EXPECT_TRUE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa2));
+  }
+
+  // Check that the non-overlapping items are not added with a wrong annotation.
+  for (uint16_t i = 0; i < 5; i++) {
+    EXPECT_FALSE(info.GetMethodHotness(MethodReference(dex1, i), psa2).IsInProfile());
+    EXPECT_FALSE(info.GetMethodHotness(MethodReference(dex2, i), psa2).IsInProfile());
+    EXPECT_FALSE(info.ContainsClass(*dex1, dex::TypeIndex(i), psa2));
+  }
+}
+
+// Verify the bulk extraction API.
+TEST_F(ProfileCompilationInfoTest, ExtractInfoWithAnnations) {
+  ProfileCompilationInfo info;
+
+  ProfileSampleAnnotation psa1("test1");
+  ProfileSampleAnnotation psa2("test2");
+
+  std::set<dex::TypeIndex> expected_classes;
+  std::set<uint16_t> expected_hot_methods;
+  std::set<uint16_t> expected_startup_methods;
+  std::set<uint16_t> expected_post_startup_methods;
+
+  for (uint16_t i = 0; i < 10; i++) {
+    ASSERT_TRUE(AddMethod(&info, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa1));
+    ASSERT_TRUE(AddClass(&info, dex1, dex::TypeIndex(i), psa1));
+    expected_hot_methods.insert(i);
+    expected_classes.insert(dex::TypeIndex(i));
+  }
+  for (uint16_t i = 5; i < 15; i++) {
+    ASSERT_TRUE(AddMethod(&info, dex1, /* method_idx= */ i, Hotness::kFlagHot, psa2));
+    ASSERT_TRUE(AddMethod(&info, dex1, /* method_idx= */ i, Hotness::kFlagStartup, psa1));
+    expected_startup_methods.insert(i);
+  }
+
+  std::set<dex::TypeIndex> classes;
+  std::set<uint16_t> hot_methods;
+  std::set<uint16_t> startup_methods;
+  std::set<uint16_t> post_startup_methods;
+
+  EXPECT_TRUE(info.GetClassesAndMethods(
+      *dex1, &classes, &hot_methods, &startup_methods, &post_startup_methods, psa1));
+  EXPECT_EQ(expected_classes, classes);
+  EXPECT_EQ(expected_hot_methods, hot_methods);
+  EXPECT_EQ(expected_startup_methods, startup_methods);
+  EXPECT_EQ(expected_post_startup_methods, post_startup_methods);
+
+  EXPECT_FALSE(info.GetClassesAndMethods(
+      *dex1,
+      &classes,
+      &hot_methods,
+      &startup_methods,
+      &post_startup_methods,
+      ProfileSampleAnnotation("new_test")));
+}
+
+// Verify the behavior for adding methods with annotations and different dex checksums.
+TEST_F(ProfileCompilationInfoTest, AddMethodsWithAnnotationAndDifferentChecksum) {
+  ProfileCompilationInfo info;
+
+  ProfileSampleAnnotation psa1("test1");
+  ProfileSampleAnnotation psa2("test2");
+
+  MethodReference ref(dex1, 0);
+  MethodReference ref_checksum_missmatch(dex1_checksum_missmatch, 1);
+
+  ASSERT_TRUE(info.AddMethod(ProfileMethodInfo(ref), Hotness::kFlagHot, psa1));
+  // Adding a method with a different dex checksum and the same annotation should fail.
+  ASSERT_FALSE(info.AddMethod(ProfileMethodInfo(ref_checksum_missmatch), Hotness::kFlagHot, psa1));
+  // However, a method with a different dex checksum and a different annotation should be ok.
+  ASSERT_TRUE(info.AddMethod(ProfileMethodInfo(ref_checksum_missmatch), Hotness::kFlagHot, psa2));
+}
+
+// Verify the behavior for searching method with annotations and different dex checksums.
+TEST_F(ProfileCompilationInfoTest, FindMethodsWithAnnotationAndDifferentChecksum) {
+  ProfileCompilationInfo info;
+
+  ProfileSampleAnnotation psa1("test1");
+
+  MethodReference ref(dex1, 0);
+  MethodReference ref_checksum_missmatch(dex1_checksum_missmatch, 0);
+
+  ASSERT_TRUE(info.AddMethod(ProfileMethodInfo(ref), Hotness::kFlagHot, psa1));
+
+  // The method should be in the profile when searched with the correct data.
+  EXPECT_TRUE(info.GetMethodHotness(ref, psa1).IsInProfile());
+  // We should get a negative result if the dex checksum  does not match.
+  EXPECT_FALSE(info.GetMethodHotness(ref_checksum_missmatch, psa1).IsInProfile());
+
+  // If we search without annotation we should have the same behaviour.
+  EXPECT_TRUE(info.GetMethodHotness(ref).IsInProfile());
+  EXPECT_FALSE(info.GetMethodHotness(ref_checksum_missmatch).IsInProfile());
+}
 }  // namespace art