Merge "Actually use --copy-and-update-profile-key argument in profman..."
diff --git a/profman/profile_assistant_test.cc b/profman/profile_assistant_test.cc
index 188d0b0..1dd1a4a 100644
--- a/profman/profile_assistant_test.cc
+++ b/profman/profile_assistant_test.cc
@@ -75,31 +75,34 @@
                     const ScratchFile& profile,
                     ProfileCompilationInfo* info,
                     uint16_t start_method_index = 0,
-                    bool reverse_dex_write_order = false) {
+                    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);
+                                      dex_location2, dex_location_checksum2,
+                                      number_of_methods1, number_of_methods2);
       Hotness::Flag flags = Hotness::kFlagPostStartup;
       if (reverse_dex_write_order) {
         ASSERT_TRUE(info->AddMethod(
-            dex_location2, dex_location_checksum2, i, kMaxMethodIds, pmi, flags));
+            dex_location2, dex_location_checksum2, i, number_of_methods2, pmi, flags));
         ASSERT_TRUE(info->AddMethod(
-            dex_location1, dex_location_checksum1, i, kMaxMethodIds, pmi, flags));
+            dex_location1, dex_location_checksum1, i, number_of_methods1, pmi, flags));
       } else {
         ASSERT_TRUE(info->AddMethod(
-            dex_location1, dex_location_checksum1, i, kMaxMethodIds, pmi, flags));
+            dex_location1, dex_location_checksum1, i, number_of_methods1, pmi, flags));
         ASSERT_TRUE(info->AddMethod(
-            dex_location2, dex_location_checksum2, i, kMaxMethodIds, pmi, flags));
+            dex_location2, dex_location_checksum2, i, number_of_methods2, pmi, flags));
       }
     }
     for (uint16_t i = 0; i < number_of_classes; i++) {
       ASSERT_TRUE(info->AddClassIndex(dex_location1,
                                       dex_location_checksum1,
                                       dex::TypeIndex(i),
-                                      kMaxMethodIds));
+                                      number_of_methods1));
     }
 
     ASSERT_TRUE(info->Save(GetFd(profile)));
@@ -143,11 +146,12 @@
 
   ProfileCompilationInfo::OfflineProfileMethodInfo GetOfflineProfileMethodInfo(
         const std::string& dex_location1, uint32_t dex_checksum1,
-        const std::string& dex_location2, uint32_t dex_checksum2) {
+        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, kMaxMethodIds);
-    pmi.dex_references.emplace_back(dex_location2, dex_checksum2, kMaxMethodIds);
+    pmi.dex_references.emplace_back(dex_location1, dex_checksum1, number_of_methods1);
+    pmi.dex_references.emplace_back(dex_location2, dex_checksum2, number_of_methods2);
 
     // Monomorphic
     for (uint16_t dex_pc = 0; dex_pc < 11; dex_pc++) {
@@ -1242,4 +1246,61 @@
   ASSERT_TRUE(expected.Equals(result));
 }
 
+TEST_F(ProfileAssistantTest, CopyAndUpdateProfileKey) {
+  ScratchFile profile1;
+  ScratchFile reference_profile;
+
+  // Use a real dex file to generate profile test data. During the copy-and-update the
+  // matching is done based on checksum so we have to match with the real thing.
+  std::vector<std::unique_ptr<const DexFile>> dex_files = OpenTestDexFiles("ProfileTestMultiDex");
+  const DexFile& d1 = *dex_files[0];
+  const DexFile& d2 = *dex_files[1];
+
+  ProfileCompilationInfo info1;
+  uint16_t num_methods_to_add = std::min(d1.NumMethodIds(), d2.NumMethodIds());
+  SetupProfile("fake-location1",
+               d1.GetLocationChecksum(),
+               "fake-location2",
+               d2.GetLocationChecksum(),
+               num_methods_to_add,
+               /*num_classes*/ 0,
+               profile1,
+               &info1,
+               /*start_method_index*/ 0,
+               /*reverse_dex_write_order*/ false,
+               /*number_of_methods1*/ d1.NumMethodIds(),
+               /*number_of_methods2*/ d2.NumMethodIds());
+
+  // Run profman and pass the dex file with --apk-fd.
+  android::base::unique_fd apk_fd(
+      open(GetTestDexFileName("ProfileTestMultiDex").c_str(), O_RDONLY));
+  ASSERT_GE(apk_fd.get(), 0);
+
+  std::string profman_cmd = GetProfmanCmd();
+  std::vector<std::string> argv_str;
+  argv_str.push_back(profman_cmd);
+  argv_str.push_back("--profile-file-fd=" + std::to_string(profile1.GetFd()));
+  argv_str.push_back("--reference-profile-file-fd=" + std::to_string(reference_profile.GetFd()));
+  argv_str.push_back("--apk-fd=" + std::to_string(apk_fd.get()));
+  argv_str.push_back("--copy-and-update-profile-key");
+  std::string error;
+
+  ASSERT_EQ(ExecAndReturnCode(argv_str, &error), 0) << error;
+
+  // Verify that we can load the result.
+  ProfileCompilationInfo result;
+  ASSERT_TRUE(reference_profile.GetFile()->ResetOffset());
+  ASSERT_TRUE(result.Load(reference_profile.GetFd()));
+
+  // Verify that the renaming was done.
+  for (uint16_t i = 0; i < num_methods_to_add; i ++) {
+      std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> pmi;
+      ASSERT_TRUE(result.GetMethod(d1.GetLocation(), d1.GetLocationChecksum(), i) != nullptr) << i;
+      ASSERT_TRUE(result.GetMethod(d2.GetLocation(), d2.GetLocationChecksum(), i) != nullptr) << i;
+
+      ASSERT_TRUE(result.GetMethod("fake-location1", d1.GetLocationChecksum(), i) == nullptr);
+      ASSERT_TRUE(result.GetMethod("fake-location2", d2.GetLocationChecksum(), i) == nullptr);
+  }
+}
+
 }  // namespace art
diff --git a/profman/profman.cc b/profman/profman.cc
index efb7fcf..5551c34 100644
--- a/profman/profman.cc
+++ b/profman/profman.cc
@@ -280,6 +280,8 @@
                         Usage);
       } else if (option.starts_with("--generate-test-profile-seed=")) {
         ParseUintOption(option, "--generate-test-profile-seed", &test_profile_seed_, Usage);
+      } else if (option.starts_with("--copy-and-update-profile-key")) {
+        copy_and_update_profile_key_ = true;
       } else {
         Usage("Unknown argument '%s'", option.data());
       }
@@ -405,9 +407,12 @@
         }
       }
     } else if (!apk_files_.empty()) {
-        if (dex_locations_.size() != apk_files_.size()) {
-            Usage("The number of apk-fds must match the number of dex-locations.");
-        }
+      if (dex_locations_.empty()) {
+        // If no dex locations are specified use the apk names as locations.
+        dex_locations_ = apk_files_;
+      } else if (dex_locations_.size() != apk_files_.size()) {
+          Usage("The number of apk-fds must match the number of dex-locations.");
+      }
     } else {
       // No APKs were specified.
       CHECK(dex_locations_.empty());
@@ -1178,7 +1183,7 @@
     return copy_and_update_profile_key_;
   }
 
-  bool CopyAndUpdateProfileKey() {
+  int32_t CopyAndUpdateProfileKey() {
     // Validate that at least one profile file was passed, as well as a reference profile.
     if (!(profile_files_.size() == 1 ^ profile_files_fd_.size() == 1)) {
       Usage("Only one profile file should be specified.");
@@ -1191,22 +1196,30 @@
       Usage("No apk files specified");
     }
 
+    static constexpr int32_t kErrorFailedToUpdateProfile = -1;
+    static constexpr int32_t kErrorFailedToSaveProfile = -2;
+    static constexpr int32_t kErrorFailedToLoadProfile = -3;
+
     bool use_fds = profile_files_fd_.size() == 1;
 
     ProfileCompilationInfo profile;
     // Do not clear if invalid. The input might be an archive.
-    if (profile.Load(profile_files_[0], /*clear_if_invalid*/ false)) {
+    bool load_ok = use_fds
+        ? profile.Load(profile_files_fd_[0])
+        : profile.Load(profile_files_[0], /*clear_if_invalid*/ false);
+    if (load_ok) {
       // Open the dex files to look up classes and methods.
       std::vector<std::unique_ptr<const DexFile>> dex_files;
       OpenApkFilesFromLocations(&dex_files);
       if (!profile.UpdateProfileKeys(dex_files)) {
-        return false;
+        return kErrorFailedToUpdateProfile;
       }
-      return use_fds
-        ? profile.Save(reference_profile_file_fd_)
-        : profile.Save(reference_profile_file_, /*bytes_written*/ nullptr);
+      bool result = use_fds
+          ? profile.Save(reference_profile_file_fd_)
+          : profile.Save(reference_profile_file_, /*bytes_written*/ nullptr);
+      return result ? 0 : kErrorFailedToSaveProfile;
     } else {
-      return false;
+      return kErrorFailedToLoadProfile;
     }
   }
 
@@ -1285,6 +1298,11 @@
   if (profman.ShouldCreateBootProfile()) {
     return profman.CreateBootProfile();
   }
+
+  if (profman.ShouldCopyAndUpdateProfileKey()) {
+    return profman.CopyAndUpdateProfileKey();
+  }
+
   // Process profile information and assess if we need to do a profile guided compilation.
   // This operation involves I/O.
   return profman.ProcessProfiles();