Enable filtering of profile data on load

This is useful for filtering profile data that does not belong to certain
apks when taking profile snapshots.

Bug: 30934496
Test: profile_compilation_info_test
Change-Id: I2100031756ef3a20eb4cbab08ddbebb067489410
diff --git a/runtime/jit/profile_compilation_info.cc b/runtime/jit/profile_compilation_info.cc
index 4bf2895..de4d02e 100644
--- a/runtime/jit/profile_compilation_info.cc
+++ b/runtime/jit/profile_compilation_info.cc
@@ -775,8 +775,14 @@
       for (; dex_classes_size > 0; dex_classes_size--) {
         uint16_t type_index;
         READ_UINT(uint16_t, buffer, type_index, error);
-        dex_pc_data->AddClass(dex_profile_index_remap.Get(dex_profile_index),
-                              dex::TypeIndex(type_index));
+        auto it = dex_profile_index_remap.find(dex_profile_index);
+        if (it == dex_profile_index_remap.end()) {
+          // If we don't have an index that's because the dex file was filtered out when loading.
+          // Set missing types on the dex pc data.
+          dex_pc_data->SetIsMissingTypes();
+        } else {
+          dex_pc_data->AddClass(it->second, dex::TypeIndex(type_index));
+        }
       }
     }
   }
@@ -1036,10 +1042,11 @@
 
 // TODO(calin): Fix this API. ProfileCompilationInfo::Load should be static and
 // return a unique pointer to a ProfileCompilationInfo upon success.
-bool ProfileCompilationInfo::Load(int fd, bool merge_classes) {
+bool ProfileCompilationInfo::Load(
+    int fd, bool merge_classes, const ProfileLoadFilterFn& filter_fn) {
   std::string error;
 
-  ProfileLoadStatus status = LoadInternal(fd, &error, merge_classes);
+  ProfileLoadStatus status = LoadInternal(fd, &error, merge_classes, filter_fn);
 
   if (status == kProfileLoadSuccess) {
     return true;
@@ -1245,7 +1252,10 @@
 
 // TODO(calin): fail fast if the dex checksums don't match.
 ProfileCompilationInfo::ProfileLoadStatus ProfileCompilationInfo::LoadInternal(
-      int32_t fd, std::string* error, bool merge_classes) {
+      int32_t fd,
+      std::string* error,
+      bool merge_classes,
+      const ProfileLoadFilterFn& filter_fn) {
   ScopedTrace trace(__PRETTY_FUNCTION__);
   DCHECK_GE(fd, 0);
 
@@ -1327,20 +1337,29 @@
   }
 
   SafeMap<uint8_t, uint8_t> dex_profile_index_remap;
-  if (!RemapProfileIndex(profile_line_headers, &dex_profile_index_remap)) {
+  if (!RemapProfileIndex(profile_line_headers, filter_fn, &dex_profile_index_remap)) {
     return kProfileLoadBadData;
   }
 
   for (uint8_t k = 0; k < number_of_dex_files; k++) {
-    // Now read the actual profile line.
-    status = ReadProfileLine(uncompressed_data,
-                             number_of_dex_files,
-                             profile_line_headers[k],
-                             dex_profile_index_remap,
-                             merge_classes,
-                             error);
-    if (status != kProfileLoadSuccess) {
-      return status;
+    if (!filter_fn(profile_line_headers[k].dex_location, profile_line_headers[k].checksum)) {
+      // We have to skip the line. Advanced the current pointer of the buffer.
+      size_t profile_line_size =
+           profile_line_headers[k].class_set_size +
+           profile_line_headers[k].method_region_size_bytes +
+           DexFileData::ComputeBitmapStorage(profile_line_headers[k].num_method_ids);
+      uncompressed_data.Advance(profile_line_size);
+    } else {
+      // Now read the actual profile line.
+      status = ReadProfileLine(uncompressed_data,
+                               number_of_dex_files,
+                               profile_line_headers[k],
+                               dex_profile_index_remap,
+                               merge_classes,
+                               error);
+      if (status != kProfileLoadSuccess) {
+        return status;
+      }
     }
   }
 
@@ -1355,12 +1374,16 @@
 
 bool ProfileCompilationInfo::RemapProfileIndex(
     const std::vector<ProfileLineHeader>& profile_line_headers,
+    const ProfileLoadFilterFn& filter_fn,
     /*out*/SafeMap<uint8_t, uint8_t>* dex_profile_index_remap) {
   // First verify that all checksums match. This will avoid adding garbage to
   // the current profile info.
   // Note that the number of elements should be very small, so this should not
   // be a performance issue.
   for (const ProfileLineHeader other_profile_line_header : profile_line_headers) {
+    if (!filter_fn(other_profile_line_header.dex_location, other_profile_line_header.checksum)) {
+      continue;
+    }
     // verify_checksum is false because we want to differentiate between a missing dex data and
     // a mismatched checksum.
     const DexFileData* dex_data = FindDexData(other_profile_line_header.dex_location,
@@ -1374,6 +1397,9 @@
   // All checksums match. Import the data.
   uint32_t num_dex_files = static_cast<uint32_t>(profile_line_headers.size());
   for (uint32_t i = 0; i < num_dex_files; i++) {
+    if (!filter_fn(profile_line_headers[i].dex_location, profile_line_headers[i].checksum)) {
+      continue;
+    }
     const DexFileData* dex_data = GetOrAddDexFileData(profile_line_headers[i].dex_location,
                                                       profile_line_headers[i].checksum,
                                                       profile_line_headers[i].num_method_ids);
@@ -2054,4 +2080,10 @@
   return true;
 }
 
+bool ProfileCompilationInfo::ProfileFilterFnAcceptAll(
+    const std::string& dex_location ATTRIBUTE_UNUSED,
+    uint32_t checksum ATTRIBUTE_UNUSED) {
+  return true;
+}
+
 }  // namespace art
diff --git a/runtime/jit/profile_compilation_info.h b/runtime/jit/profile_compilation_info.h
index 350ce9e..1973f3f 100644
--- a/runtime/jit/profile_compilation_info.h
+++ b/runtime/jit/profile_compilation_info.h
@@ -306,7 +306,19 @@
   // Load or Merge profile information from the given file descriptor.
   // If the current profile is non-empty the load will fail.
   // If merge_classes is set to false, classes will not be merged/loaded.
-  bool Load(int fd, bool merge_classes = true);
+  // If filter_fn is present, it will be used to filter out profile data belonging
+  // to dex file which do not comply with the filter
+  // (i.e. for which filter_fn(dex_location, dex_checksum) is false).
+  using ProfileLoadFilterFn = std::function<bool(const std::string&, uint32_t)>;
+  // Profile filter method which accepts all dex locations.
+  // This is convenient to use when we need to accept all locations without repeating the same
+  // lambda.
+  static bool ProfileFilterFnAcceptAll(const std::string& dex_location, uint32_t checksum);
+
+  bool Load(
+      int fd,
+      bool merge_classes = true,
+      const ProfileLoadFilterFn& filter_fn = ProfileFilterFnAcceptAll);
 
   // Verify integrity of the profile file with the provided dex files.
   // If there exists a DexData object which maps to a dex_file, then it verifies that:
@@ -459,14 +471,21 @@
           class_set(std::less<dex::TypeIndex>(), allocator->Adapter(kArenaAllocProfile)),
           num_method_ids(num_methods),
           bitmap_storage(allocator->Adapter(kArenaAllocProfile)) {
-      const size_t num_bits = num_method_ids * kBitmapIndexCount;
-      bitmap_storage.resize(RoundUp(num_bits, kBitsPerByte) / kBitsPerByte);
+      bitmap_storage.resize(ComputeBitmapStorage(num_method_ids));
       if (!bitmap_storage.empty()) {
         method_bitmap =
-            BitMemoryRegion(MemoryRegion(&bitmap_storage[0], bitmap_storage.size()), 0, num_bits);
+            BitMemoryRegion(MemoryRegion(
+                &bitmap_storage[0], bitmap_storage.size()), 0, ComputeBitmapBits(num_method_ids));
       }
     }
 
+    static size_t ComputeBitmapBits(uint32_t num_method_ids) {
+      return num_method_ids * kBitmapIndexCount;
+    }
+    static size_t ComputeBitmapStorage(uint32_t num_method_ids) {
+      return RoundUp(ComputeBitmapBits(num_method_ids), kBitsPerByte) / kBitsPerByte;
+    }
+
     bool operator==(const DexFileData& other) const {
       return checksum == other.checksum && method_map == other.method_map;
     }
@@ -689,10 +708,12 @@
                                /*out*/ std::unique_ptr<ProfileSource>* source,
                                /*out*/ std::string* error);
 
-  // Entry point for profile loding functionality.
-  ProfileLoadStatus LoadInternal(int32_t fd,
-                                 std::string* error,
-                                 bool merge_classes = true);
+  // Entry point for profile loading functionality.
+  ProfileLoadStatus LoadInternal(
+      int32_t fd,
+      std::string* error,
+      bool merge_classes = true,
+      const ProfileLoadFilterFn& filter_fn = ProfileFilterFnAcceptAll);
 
   // Read the profile header from the given fd and store the number of profile
   // lines into number_of_dex_files.
@@ -736,6 +757,7 @@
   // The method generates mapping of profile indices while merging a new profile
   // data into current data. It returns true, if the mapping was successful.
   bool RemapProfileIndex(const std::vector<ProfileLineHeader>& profile_line_headers,
+                         const ProfileLoadFilterFn& filter_fn,
                          /*out*/SafeMap<uint8_t, uint8_t>* dex_profile_index_remap);
 
   // Read the inline cache encoding from line_bufer into inline_cache.
diff --git a/runtime/jit/profile_compilation_info_test.cc b/runtime/jit/profile_compilation_info_test.cc
index b4265d1..4ac11ee 100644
--- a/runtime/jit/profile_compilation_info_test.cc
+++ b/runtime/jit/profile_compilation_info_test.cc
@@ -314,6 +314,10 @@
     }
   }
 
+  bool IsEmpty(const ProfileCompilationInfo& info) {
+    return info.IsEmpty();
+  }
+
   // Cannot sizeof the actual arrays so hard code the values here.
   // They should not change anyway.
   static constexpr int kProfileMagicSize = 4;
@@ -1124,4 +1128,169 @@
   ASSERT_FALSE(info.UpdateProfileKeys(dex_files));
 }
 
+TEST_F(ProfileCompilationInfoTest, FilteredLoading) {
+  ScratchFile profile;
+
+  ProfileCompilationInfo saved_info;
+  ProfileCompilationInfo::OfflineProfileMethodInfo pmi = GetOfflineProfileMethodInfo();
+
+  // Add methods with inline caches.
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    // Add a method which is part of the same dex file as one of the class from the inline caches.
+    ASSERT_TRUE(AddMethod("dex_location1", /* checksum */ 1, method_idx, pmi, &saved_info));
+    ASSERT_TRUE(AddMethod("dex_location2", /* checksum */ 2, method_idx, pmi, &saved_info));
+    // Add a method which is outside the set of dex files.
+    ASSERT_TRUE(AddMethod("dex_location4", /* checksum */ 4, method_idx, pmi, &saved_info));
+  }
+
+  ASSERT_TRUE(saved_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());
+
+  // Filter out dex locations. Keep only dex_location1 and dex_location2.
+  ProfileCompilationInfo::ProfileLoadFilterFn filter_fn =
+      [](const std::string& dex_location, uint32_t checksum) -> bool {
+          return (dex_location == "dex_location1" && checksum == 1)
+              || (dex_location == "dex_location3" && checksum == 3);
+        };
+  ASSERT_TRUE(loaded_info.Load(GetFd(profile), true, filter_fn));
+
+  // Verify that we filtered out locations during load.
+
+  // Dex location 2 and 4 should have been filtered out
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    ASSERT_TRUE(nullptr == loaded_info.GetMethod("dex_location2", /* checksum */ 2, method_idx));
+    ASSERT_TRUE(nullptr == loaded_info.GetMethod("dex_location4", /* checksum */ 4, method_idx));
+  }
+
+  // Dex location 1 should have all all the inline caches referencing dex location 2 set to
+  // missing types.
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    // The methods for dex location 1 should be in the profile data.
+    std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> loaded_pmi1 =
+      loaded_info.GetMethod("dex_location1", /* checksum */ 1, /* method_idx */ method_idx);
+    ASSERT_TRUE(loaded_pmi1 != nullptr);
+
+    // Verify the inline cache.
+    // Everything should be as constructed by GetOfflineProfileMethodInfo with the exception
+    // of the inline caches referring types from dex_location2.
+    // These should be set to IsMissingType.
+    ProfileCompilationInfo::InlineCacheMap* ic_map = CreateInlineCacheMap();
+
+    // Monomorphic types should remain the same as dex_location1 was kept.
+    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);
+    }
+    // Polymorphic inline cache should have been transformed to IsMissingType due to
+    // the removal of dex_location2.
+    for (uint16_t dex_pc = 11; dex_pc < 22; dex_pc++) {
+      ProfileCompilationInfo::DexPcData dex_pc_data(allocator_.get());
+      dex_pc_data.SetIsMissingTypes();
+      ic_map->Put(dex_pc, dex_pc_data);
+    }
+
+    // Megamorphic are not affected by removal of dex files.
+    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);
+    }
+    // Missing types are not affected be removal of dex files.
+    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);
+    }
+
+    ProfileCompilationInfo::OfflineProfileMethodInfo expected_pmi(ic_map);
+
+    // The dex references should not have  dex_location2 in the list.
+    expected_pmi.dex_references.emplace_back("dex_location1", /* checksum */1, kMaxMethodIds);
+    expected_pmi.dex_references.emplace_back("dex_location3", /* checksum */3, kMaxMethodIds);
+
+    // Now check that we get back what we expect.
+    ASSERT_TRUE(*loaded_pmi1 == expected_pmi);
+  }
+}
+
+TEST_F(ProfileCompilationInfoTest, FilteredLoadingRemoveAll) {
+  ScratchFile profile;
+
+  ProfileCompilationInfo saved_info;
+  ProfileCompilationInfo::OfflineProfileMethodInfo pmi = GetOfflineProfileMethodInfo();
+
+  // Add methods with inline caches.
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    // Add a method which is part of the same dex file as one of the class from the inline caches.
+    ASSERT_TRUE(AddMethod("dex_location1", /* checksum */ 1, method_idx, pmi, &saved_info));
+    ASSERT_TRUE(AddMethod("dex_location2", /* checksum */ 2, method_idx, pmi, &saved_info));
+    // Add a method which is outside the set of dex files.
+    ASSERT_TRUE(AddMethod("dex_location4", /* checksum */ 4, method_idx, pmi, &saved_info));
+  }
+
+  ASSERT_TRUE(saved_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());
+
+  // Remove all elements.
+  ProfileCompilationInfo::ProfileLoadFilterFn filter_fn =
+      [](const std::string&, uint32_t) -> bool { return false; };
+  ASSERT_TRUE(loaded_info.Load(GetFd(profile), true, filter_fn));
+
+  // Verify that we filtered out everything.
+  ASSERT_TRUE(IsEmpty(loaded_info));
+}
+
+TEST_F(ProfileCompilationInfoTest, FilteredLoadingKeepAll) {
+  ScratchFile profile;
+
+  ProfileCompilationInfo saved_info;
+  ProfileCompilationInfo::OfflineProfileMethodInfo pmi = GetOfflineProfileMethodInfo();
+
+  // Add methods with inline caches.
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    // Add a method which is part of the same dex file as one of the
+    // class from the inline caches.
+    ASSERT_TRUE(AddMethod("dex_location1", /* checksum */ 1, method_idx, pmi, &saved_info));
+    // Add a method which is outside the set of dex files.
+    ASSERT_TRUE(AddMethod("dex_location4", /* checksum */ 4, method_idx, pmi, &saved_info));
+  }
+
+  ASSERT_TRUE(saved_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());
+
+  // Keep all elements.
+  ProfileCompilationInfo::ProfileLoadFilterFn filter_fn =
+      [](const std::string&, uint32_t) -> bool { return true; };
+  ASSERT_TRUE(loaded_info.Load(GetFd(profile), true, filter_fn));
+
+
+  ASSERT_TRUE(loaded_info.Equals(saved_info));
+
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> loaded_pmi1 =
+        loaded_info.GetMethod("dex_location1", /* checksum */ 1, method_idx);
+    ASSERT_TRUE(loaded_pmi1 != nullptr);
+    ASSERT_TRUE(*loaded_pmi1 == pmi);
+  }
+  for (uint16_t method_idx = 0; method_idx < 10; method_idx++) {
+    std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> loaded_pmi2 =
+        loaded_info.GetMethod("dex_location4", /* checksum */ 4, method_idx);
+    ASSERT_TRUE(loaded_pmi2 != nullptr);
+    ASSERT_TRUE(*loaded_pmi2 == pmi);
+  }
+}
+
 }  // namespace art