Avoid scanning pre-installed APEX twice during boot sequence

Now that we can retrieve ApexFiles from ApexFileRepoitory we don't
have to scan apex files in ScanAndGroupApexFiles, which is called
on every boot.

Bug: 179497746
Test: atest ApexTestCases
Change-Id: Ib70a38c8b35bdb82dd6b31b6dd9662f9d815de1b
diff --git a/apexd/apex_file_repository.cpp b/apexd/apex_file_repository.cpp
index 8b11fb3..9de847c 100644
--- a/apexd/apex_file_repository.cpp
+++ b/apexd/apex_file_repository.cpp
@@ -51,6 +51,7 @@
     return all_apex_files.error();
   }
 
+  // TODO(b/179248390): scan parallelly if possible
   for (const auto& file : *all_apex_files) {
     Result<ApexFile> apex_file = ApexFile::Open(file);
     if (!apex_file.ok()) {
@@ -113,10 +114,12 @@
     return all_apex_files.error();
   }
 
+  // TODO(b/179248390): scan parallelly if possible
   for (const auto& file : *all_apex_files) {
     Result<ApexFile> apex_file = ApexFile::Open(file);
     if (!apex_file.ok()) {
-      return Error() << "Failed to open " << file << " : " << apex_file.error();
+      LOG(ERROR) << "Failed to open " << file << " : " << apex_file.error();
+      continue;
     }
 
     const std::string& name = apex_file->GetManifest().name();
@@ -199,7 +202,7 @@
 }
 
 std::vector<std::reference_wrapper<const ApexFile>>
-ApexFileRepository::GetPreInstalledApexFiles() {
+ApexFileRepository::GetPreInstalledApexFiles() const {
   std::vector<std::reference_wrapper<const ApexFile>> result;
   for (const auto& it : pre_installed_store_) {
     result.emplace_back(std::cref(it.second));
@@ -208,7 +211,7 @@
 }
 
 std::vector<std::reference_wrapper<const ApexFile>>
-ApexFileRepository::GetDataApexFiles() {
+ApexFileRepository::GetDataApexFiles() const {
   std::vector<std::reference_wrapper<const ApexFile>> result;
   for (const auto& it : data_store_) {
     result.emplace_back(std::cref(it.second));
@@ -216,5 +219,35 @@
   return std::move(result);
 }
 
+// Group pre-installed APEX and data APEX by name
+std::unordered_map<std::string,
+                   std::vector<std::reference_wrapper<const ApexFile>>>
+ApexFileRepository::AllApexFilesByName() const {
+  // Collect all apex files
+  std::vector<std::reference_wrapper<const ApexFile>> all_apex_files;
+  auto pre_installed_apexs = GetPreInstalledApexFiles();
+  auto data_apexs = GetDataApexFiles();
+  std::move(pre_installed_apexs.begin(), pre_installed_apexs.end(),
+            std::back_inserter(all_apex_files));
+  std::move(data_apexs.begin(), data_apexs.end(),
+            std::back_inserter(all_apex_files));
+
+  // Group them by name
+  std::unordered_map<std::string,
+                     std::vector<std::reference_wrapper<const ApexFile>>>
+      result;
+  for (const auto& apex_file_ref : all_apex_files) {
+    const ApexFile& apex_file = apex_file_ref.get();
+    const std::string& package_name = apex_file.GetManifest().name();
+    if (result.find(package_name) == result.end()) {
+      result[package_name] =
+          std::vector<std::reference_wrapper<const ApexFile>>{};
+    }
+    result[package_name].emplace_back(apex_file_ref);
+  }
+
+  return std::move(result);
+}
+
 }  // namespace apex
 }  // namespace android
diff --git a/apexd/apex_file_repository.h b/apexd/apex_file_repository.h
index e28fa51..6d3199b 100644
--- a/apexd/apex_file_repository.h
+++ b/apexd/apex_file_repository.h
@@ -83,11 +83,16 @@
   bool IsDecompressedApex(const ApexFile& apex) const;
 
   // Returns reference to all pre-installed APEX on device
-  std::vector<std::reference_wrapper<const ApexFile>>
-  GetPreInstalledApexFiles();
+  std::vector<std::reference_wrapper<const ApexFile>> GetPreInstalledApexFiles()
+      const;
 
   // Returns reference to all data APEX on device
-  std::vector<std::reference_wrapper<const ApexFile>> GetDataApexFiles();
+  std::vector<std::reference_wrapper<const ApexFile>> GetDataApexFiles() const;
+
+  // Group all ApexFiles on device by their package name
+  std::unordered_map<std::string,
+                     std::vector<std::reference_wrapper<const ApexFile>>>
+  AllApexFilesByName() const;
 
  private:
   // Non-copyable && non-moveable.
diff --git a/apexd/apex_file_repository_test.cpp b/apexd/apex_file_repository_test.cpp
index 923d681..0e9000b 100644
--- a/apexd/apex_file_repository_test.cpp
+++ b/apexd/apex_file_repository_test.cpp
@@ -420,5 +420,40 @@
                                    ApexFileEq(ByRef(*pre_apex_2))));
 }
 
+TEST(ApexdUnitTest, AllApexFilesByName) {
+  TemporaryDir built_in_dir;
+  fs::copy(GetTestFile("apex.apexd_test.apex"), built_in_dir.path);
+  fs::copy(GetTestFile("com.android.apex.cts.shim.apex"), built_in_dir.path);
+  fs::copy(GetTestFile("com.android.apex.compressed.v1.capex"),
+           built_in_dir.path);
+  ApexFileRepository instance;
+  ASSERT_TRUE(IsOk(instance.AddPreInstalledApex({built_in_dir.path})));
+
+  TemporaryDir data_dir;
+  fs::copy(GetTestFile("com.android.apex.cts.shim.v2.apex"), data_dir.path);
+  ASSERT_TRUE(IsOk(instance.AddDataApex(data_dir.path)));
+
+  auto result = instance.AllApexFilesByName();
+
+  // Verify the contents of result
+  auto apexd_test_file = ApexFile::Open(
+      StringPrintf("%s/apex.apexd_test.apex", built_in_dir.path));
+  auto shim_v1 = ApexFile::Open(
+      StringPrintf("%s/com.android.apex.cts.shim.apex", built_in_dir.path));
+  auto compressed_apex = ApexFile::Open(StringPrintf(
+      "%s/com.android.apex.compressed.v1.capex", built_in_dir.path));
+  auto shim_v2 = ApexFile::Open(
+      StringPrintf("%s/com.android.apex.cts.shim.v2.apex", data_dir.path));
+
+  ASSERT_EQ(result.size(), 3u);
+  ASSERT_THAT(result[apexd_test_file->GetManifest().name()],
+              UnorderedElementsAre(ApexFileEq(ByRef(*apexd_test_file))));
+  ASSERT_THAT(result[shim_v1->GetManifest().name()],
+              UnorderedElementsAre(ApexFileEq(ByRef(*shim_v1)),
+                                   ApexFileEq(ByRef(*shim_v2))));
+  ASSERT_THAT(result[compressed_apex->GetManifest().name()],
+              UnorderedElementsAre(ApexFileEq(ByRef(*compressed_apex))));
+}
+
 }  // namespace apex
 }  // namespace android
diff --git a/apexd/apexd.cpp b/apexd/apexd.cpp
index 6ae38bc..cc82663 100644
--- a/apexd/apexd.cpp
+++ b/apexd/apexd.cpp
@@ -61,6 +61,7 @@
 #include <sys/sysinfo.h>
 #include <sys/types.h>
 #include <unistd.h>
+#include <algorithm>
 
 #include <algorithm>
 #include <array>
@@ -1458,6 +1459,7 @@
 
 namespace {
 
+// TODO(b/179497746): Avoid scanning apex directly here
 Result<std::vector<ApexFile>> ScanApexFiles(const char* apex_package_dir,
                                             bool include_compressed = false) {
   LOG(INFO) << "Scanning " << apex_package_dir << " looking for APEX packages.";
@@ -1512,13 +1514,14 @@
   return ret;
 }
 
-Result<void> ActivateApexPackages(const std::vector<ApexFile>& apexes) {
+Result<void> ActivateApexPackages(
+    const std::vector<std::reference_wrapper<const ApexFile>>& apexes) {
   const auto& packages_with_code = GetActivePackagesMap();
   size_t skipped_cnt = 0;
   std::queue<const ApexFile*> apex_queue;
   std::mutex apex_queue_mutex;
 
-  for (const auto& apex : apexes) {
+  for (const ApexFile& apex : apexes) {
     if (!apex.GetManifest().providesharedapexlibs()) {
       uint64_t new_version =
           static_cast<uint64_t>(apex.GetManifest().version());
@@ -1578,12 +1581,16 @@
 
 }  // namespace
 
+// TODO(b/179497746): Avoid scanning APEX here
 Result<void> ScanPackagesDirAndActivate(const char* apex_package_dir) {
   auto apexes = ScanApexFiles(apex_package_dir);
   if (!apexes.ok()) {
     return apexes.error();
   }
-  return ActivateApexPackages(*apexes);
+  std::vector<std::reference_wrapper<const ApexFile>> apexes_ref;
+  std::transform(apexes->begin(), apexes->end(), std::back_inserter(apexes_ref),
+                 [](const auto& x) { return std::cref(x); });
+  return ActivateApexPackages(apexes_ref);
 }
 
 /**
@@ -1837,6 +1844,7 @@
   BootCompletedCleanup();
 }
 
+// Returns true if any session gets staged
 void ScanStagedSessionsDirAndStage() {
   LOG(INFO) << "Scanning " << ApexSession::GetSessionsDir()
             << " looking for sessions to be activated.";
@@ -2241,7 +2249,11 @@
   }
 
   // Now activate bootstrap apexes.
-  if (auto ret = ActivateApexPackages(bootstrap_apexes); !ret.ok()) {
+  std::vector<std::reference_wrapper<const ApexFile>> bootstrap_apexes_ref;
+  std::transform(bootstrap_apexes.begin(), bootstrap_apexes.end(),
+                 std::back_inserter(bootstrap_apexes_ref),
+                 [](const auto& x) { return std::cref(x); });
+  if (auto ret = ActivateApexPackages(bootstrap_apexes_ref); !ret.ok()) {
     LOG(ERROR) << "Failed to activate bootstrap apex files : " << ret.error();
     return 1;
   }
@@ -2290,43 +2302,9 @@
                << status.error();
     return;
   }
-  status = instance.AddDataApex(kActiveApexPackagesDataDir);
-  if (!status.ok()) {
-    LOG(ERROR) << "Failed to collect data APEX files : " << status.error();
-    return;
-  }
-
   gMountedApexes.PopulateFromMounts();
 }
 
-// Scans all APEX in the given directories and groups them by their package name
-std::unordered_map<std::string, std::vector<ApexFile>> ScanAndGroupApexFiles(
-    const std::vector<std::string>& dirs_to_scan) {
-  LOG(INFO) << "Scanning all apex";
-  std::unordered_map<std::string, std::vector<ApexFile>> result;
-
-  // TODO(b/179248390): scan parallelly if possible
-  for (const auto& dir : dirs_to_scan) {
-    auto scan_status = ScanApexFiles(dir.c_str(), true /*include_compressed*/);
-    if (scan_status.ok()) {
-      for (ApexFile& apex_file : *scan_status) {
-        const std::string& package_name = apex_file.GetManifest().name();
-        if (result.find(package_name) == result.end()) {
-          result[package_name] = std::vector<ApexFile>{};
-        }
-        result[package_name].emplace_back(std::move(apex_file));
-      }
-    } else {
-      LOG(ERROR) << "Failed to scan APEX packages from " << dir << " : "
-                 << scan_status.error();
-      if (auto revert = RevertActiveSessionsAndReboot(""); !revert.ok()) {
-        LOG(ERROR) << "Failed to revert : " << revert.error();
-      }
-    }
-  }
-  return result;
-}
-
 /**
  * For every package X, there can be at most two APEX, pre-installed vs
  * installed on data. We usually select only one of these APEX for each package
@@ -2342,15 +2320,18 @@
  * @param all_apex all the APEX grouped by their package name
  * @return list of ApexFile that needs to be activated
  */
-std::vector<ApexFile> SelectApexForActivation(
-    std::unordered_map<std::string, std::vector<ApexFile>>&& all_apex,
+std::vector<std::reference_wrapper<const ApexFile>> SelectApexForActivation(
+    const std::unordered_map<
+        std::string, std::vector<std::reference_wrapper<const ApexFile>>>&
+        all_apex,
     const ApexFileRepository& instance) {
   LOG(INFO) << "Selecting APEX for activation";
-  std::vector<ApexFile> activation_list;
+  std::vector<std::reference_wrapper<const ApexFile>> activation_list;
   // For every package X, select which APEX to activate
   for (auto& apex_it : all_apex) {
     const std::string& package_name = apex_it.first;
-    std::vector<ApexFile> apex_files = std::move(apex_it.second);
+    const std::vector<std::reference_wrapper<const ApexFile>>& apex_files =
+        apex_it.second;
 
     if (apex_files.size() > 2 || apex_files.size() == 0) {
       LOG(FATAL) << "Unexpectedly found more than two versions or none for "
@@ -2368,15 +2349,20 @@
 
     if (apex_files.size() == 1) {
       LOG(DEBUG) << "Selecting the only APEX: " << package_name << " "
-                 << apex_files[0].GetPath();
-      activation_list.emplace_back(std::move(apex_files[0]));
+                 << apex_files[0].get().GetPath();
+      activation_list.emplace_back(apex_files[0]);
       continue;
     }
 
+    // TODO(b/179497746): Now that we are dealing with list of reference, this
+    //  selection process can be simplified by sorting the vector.
+
     // Given an APEX A and the version of the other APEX B, should we activate
     // it?
     auto select_apex = [&instance, &activation_list](
-                           ApexFile&& a, const int version_b) mutable {
+                           const std::reference_wrapper<const ApexFile>& a_ref,
+                           const int version_b) mutable {
+      const ApexFile& a = a_ref.get();
       // APEX that provides shared library always gets activated
       const bool provides_shared_apex_libs =
           a.GetManifest().providesharedapexlibs();
@@ -2393,13 +2379,13 @@
           same_version_priority_to_data || decompressed) {
         LOG(DEBUG) << "Selecting between two APEX: " << a.GetManifest().name()
                    << " " << a.GetPath();
-        activation_list.emplace_back(std::move(a));
+        activation_list.emplace_back(a_ref);
       }
     };
-    const int version_0 = apex_files[0].GetManifest().version();
-    const int version_1 = apex_files[1].GetManifest().version();
-    select_apex(std::move(apex_files[0]), version_1);
-    select_apex(std::move(apex_files[1]), version_0);
+    const int version_0 = apex_files[0].get().GetManifest().version();
+    const int version_1 = apex_files[1].get().GetManifest().version();
+    select_apex(apex_files[0].get(), version_1);
+    select_apex(apex_files[1].get(), version_0);
   }
   return activation_list;
 }
@@ -2409,9 +2395,8 @@
  * link it to kActiveApexPackagesDataDir. Returns list of decompressed APEX.
  */
 std::vector<ApexFile> ProcessCompressedApex(
-    std::vector<ApexFile>&& compressed_apex,
-    const std::string& decompression_dir = kApexDecompressedDir,
-    const std::string& active_apex_dir = kActiveApexPackagesDataDir) {
+    const std::vector<std::reference_wrapper<const ApexFile>>& compressed_apex,
+    const std::string& decompression_dir, const std::string& active_apex_dir) {
   LOG(INFO) << "Processing compressed APEX";
   std::vector<ApexFile> decompressed_apex_list;
   for (const ApexFile& apex_file : compressed_apex) {
@@ -2511,41 +2496,45 @@
   // If there is any new apex to be installed on /data/app-staging, hardlink
   // them to /data/apex/active first.
   ScanStagedSessionsDirAndStage();
+  if (auto status = ApexFileRepository::GetInstance().AddDataApex(
+          kActiveApexPackagesDataDir);
+      !status.ok()) {
+    LOG(ERROR) << "Failed to collect data APEX files : " << status.error();
+  }
+
   auto status = ResumeRevertIfNeeded();
   if (!status.ok()) {
     LOG(ERROR) << "Failed to resume revert : " << status.error();
   }
 
-  // Scan every APEX on device
-  std::vector<std::string> dirs_to_scan = kApexPackageBuiltinDirs;
-  dirs_to_scan.push_back(kActiveApexPackagesDataDir);
-  std::unordered_map<std::string, std::vector<ApexFile>> all_apex =
-      ScanAndGroupApexFiles(dirs_to_scan);
+  // Group every ApexFile on device by name
+  const auto& instance = ApexFileRepository::GetInstance();
+  const auto& all_apex = instance.AllApexFilesByName();
   // There can be multiple APEX packages with package name X. Determine which
   // one to activate.
-  const auto& instance = ApexFileRepository::GetInstance();
-  std::vector<ApexFile> activation_list =
-      SelectApexForActivation(std::move(all_apex), instance);
+  auto activation_list = SelectApexForActivation(all_apex, instance);
 
   // Process compressed APEX, if any
-  std::vector<ApexFile> compressed_apex;
+  std::vector<std::reference_wrapper<const ApexFile>> compressed_apex;
   for (auto it = activation_list.begin(); it != activation_list.end();) {
-    if (it->IsCompressed()) {
-      compressed_apex.emplace_back(std::move(*it));
+    if (it->get().IsCompressed()) {
+      compressed_apex.emplace_back(*it);
       it = activation_list.erase(it);
     } else {
       it++;
     }
   }
+  std::vector<ApexFile> decompressed_apex;
   if (!compressed_apex.empty()) {
-    auto decompressed_apex = ProcessCompressedApex(std::move(compressed_apex));
-    std::move(decompressed_apex.begin(), decompressed_apex.end(),
-              std::back_inserter(activation_list));
+    decompressed_apex = ProcessCompressedApex(compressed_apex);
+    for (const ApexFile& apex_file : decompressed_apex) {
+      activation_list.emplace_back(std::cref(apex_file));
+    }
   }
 
   int data_apex_cnt = std::count_if(
-      activation_list.begin(), activation_list.end(), [](const ApexFile& a) {
-        return !ApexFileRepository::GetInstance().IsPreInstalledApex(a);
+      activation_list.begin(), activation_list.end(), [](const auto& a) {
+        return !ApexFileRepository::GetInstance().IsPreInstalledApex(a.get());
       });
   if (data_apex_cnt > 0) {
     Result<void> pre_allocate = loop::PreAllocateLoopDevices(data_apex_cnt);
diff --git a/apexd/apexd.h b/apexd/apexd.h
index 5c6186d..5144f82 100644
--- a/apexd/apexd.h
+++ b/apexd/apexd.h
@@ -104,18 +104,18 @@
 // Must only be called during boot (i.e apexd.status is not "ready" or
 // "activated").
 void OnStart();
-// Scans all APEX in the given directories and groups them by their package name
-std::unordered_map<std::string, std::vector<ApexFile>> ScanAndGroupApexFiles(
-    const std::vector<std::string>& dirs_to_scan);
 // For every package X, there can be at most two APEX, pre-installed vs
 // installed on data. We decide which ones should be activated and return them
 // as a list
-std::vector<ApexFile> SelectApexForActivation(
-    std::unordered_map<std::string, std::vector<ApexFile>>&& all_apex,
+std::vector<std::reference_wrapper<const ApexFile>> SelectApexForActivation(
+    const std::unordered_map<
+        std::string, std::vector<std::reference_wrapper<const ApexFile>>>&
+        all_apex,
     const ApexFileRepository& instance);
 std::vector<ApexFile> ProcessCompressedApex(
-    std::vector<ApexFile>&& compressed_apex,
-    const std::string& decompression_dir, const std::string& active_apex_dir);
+    const std::vector<std::reference_wrapper<const ApexFile>>& compressed_apex,
+    const std::string& decompression_dir = kApexDecompressedDir,
+    const std::string& active_apex_dir = kActiveApexPackagesDataDir);
 // Notifies system that apexes are activated by setting apexd.status property to
 // "activated".
 // Must only be called during boot (i.e. apexd.status is not "ready" or
diff --git a/apexd/apexd_test.cpp b/apexd/apexd_test.cpp
index 4c6ea52..8205f14 100644
--- a/apexd/apexd_test.cpp
+++ b/apexd/apexd_test.cpp
@@ -44,39 +44,6 @@
   return GetTestDataDir() + "/" + name;
 }
 
-TEST(ApexdUnitTest, ScanAndGroupApexFiles) {
-  TemporaryDir built_in_dir;
-  fs::copy(GetTestFile("apex.apexd_test.apex"), built_in_dir.path);
-  fs::copy(GetTestFile("com.android.apex.cts.shim.apex"), built_in_dir.path);
-  fs::copy(GetTestFile("com.android.apex.compressed.v1.capex"),
-           built_in_dir.path);
-
-  TemporaryDir data_dir;
-  fs::copy(GetTestFile("com.android.apex.cts.shim.v2.apex"), data_dir.path);
-
-  std::vector<std::string> dirs_to_scan{built_in_dir.path, data_dir.path};
-  auto result = ScanAndGroupApexFiles(dirs_to_scan);
-
-  // Verify the contents of result
-  auto apexd_test_file = ApexFile::Open(
-      StringPrintf("%s/apex.apexd_test.apex", built_in_dir.path));
-  auto shim_v1 = ApexFile::Open(
-      StringPrintf("%s/com.android.apex.cts.shim.apex", built_in_dir.path));
-  auto compressed_apex = ApexFile::Open(StringPrintf(
-      "%s/com.android.apex.compressed.v1.capex", built_in_dir.path));
-  auto shim_v2 = ApexFile::Open(
-      StringPrintf("%s/com.android.apex.cts.shim.v2.apex", data_dir.path));
-
-  ASSERT_EQ(result.size(), 3u);
-  ASSERT_THAT(result[apexd_test_file->GetManifest().name()],
-              UnorderedElementsAre(ApexFileEq(ByRef(*apexd_test_file))));
-  ASSERT_THAT(result[shim_v1->GetManifest().name()],
-              UnorderedElementsAre(ApexFileEq(ByRef(*shim_v1)),
-                                   ApexFileEq(ByRef(*shim_v2))));
-  ASSERT_THAT(result[compressed_apex->GetManifest().name()],
-              UnorderedElementsAre(ApexFileEq(ByRef(*compressed_apex))));
-}
-
 // Apex that does not have pre-installed version, does not get selected
 TEST(ApexdUnitTest, ApexMustHavePreInstalledVersionForSelection) {
   TemporaryDir built_in_dir;
@@ -85,29 +52,41 @@
   fs::copy(
       GetTestFile("com.android.apex.test.sharedlibs_generated.v1.libvX.apex"),
       built_in_dir.path);
-
-  // Pre-installed information is not initialized
   ApexFileRepository instance;
-  std::vector<std::string> dirs_to_scan{built_in_dir.path};
-  auto all_apex = ScanAndGroupApexFiles(dirs_to_scan);
-  auto result = SelectApexForActivation(std::move(all_apex), instance);
-  ASSERT_EQ(result.size(), 0u);
-
-  // Once initialized, pre-installed APEX should get selected
+  // Pre-installed data needs to be present so that we can add data apex
   ASSERT_TRUE(IsOk(instance.AddPreInstalledApex({built_in_dir.path})));
-  all_apex = ScanAndGroupApexFiles(dirs_to_scan);
-  result = SelectApexForActivation(std::move(all_apex), instance);
-  ASSERT_EQ(result.size(), 3u);
-  auto apexd_test_file = ApexFile::Open(
-      StringPrintf("%s/apex.apexd_test.apex", built_in_dir.path));
+
+  TemporaryDir data_dir;
+  fs::copy(GetTestFile("apex.apexd_test.apex"), data_dir.path);
+  fs::copy(GetTestFile("com.android.apex.cts.shim.apex"), data_dir.path);
+  fs::copy(
+      GetTestFile("com.android.apex.test.sharedlibs_generated.v1.libvX.apex"),
+      data_dir.path);
+  ASSERT_TRUE(IsOk(instance.AddDataApex(data_dir.path)));
+
+  const auto all_apex = instance.AllApexFilesByName();
+  // Pass a blank instance so that the data apex files are not considered
+  // pre-installed
+  const ApexFileRepository instance_blank;
+  auto result = SelectApexForActivation(all_apex, instance_blank);
+  ASSERT_EQ(result.size(), 0u);
+  // When passed proper instance they should get selected
+  result = SelectApexForActivation(all_apex, instance);
+  ASSERT_EQ(result.size(), 4u);
+  auto apexd_test_file =
+      ApexFile::Open(StringPrintf("%s/apex.apexd_test.apex", data_dir.path));
   auto shim_v1 = ApexFile::Open(
-      StringPrintf("%s/com.android.apex.cts.shim.apex", built_in_dir.path));
-  auto shared_lib = ApexFile::Open(StringPrintf(
+      StringPrintf("%s/com.android.apex.cts.shim.apex", data_dir.path));
+  auto shared_lib_1 = ApexFile::Open(StringPrintf(
       "%s/com.android.apex.test.sharedlibs_generated.v1.libvX.apex",
       built_in_dir.path));
+  auto shared_lib_2 = ApexFile::Open(StringPrintf(
+      "%s/com.android.apex.test.sharedlibs_generated.v1.libvX.apex",
+      data_dir.path));
   ASSERT_THAT(result, UnorderedElementsAre(ApexFileEq(ByRef(*apexd_test_file)),
                                            ApexFileEq(ByRef(*shim_v1)),
-                                           ApexFileEq(ByRef(*shared_lib))));
+                                           ApexFileEq(ByRef(*shared_lib_1)),
+                                           ApexFileEq(ByRef(*shared_lib_2))));
 }
 
 // Higher version gets priority when selecting for activation
@@ -115,18 +94,16 @@
   TemporaryDir built_in_dir;
   fs::copy(GetTestFile("apex.apexd_test_v2.apex"), built_in_dir.path);
   fs::copy(GetTestFile("com.android.apex.cts.shim.apex"), built_in_dir.path);
-
-  // Initialize pre-installed APEX information
   ApexFileRepository instance;
   ASSERT_TRUE(IsOk(instance.AddPreInstalledApex({built_in_dir.path})));
 
   TemporaryDir data_dir;
   fs::copy(GetTestFile("apex.apexd_test.apex"), data_dir.path);
   fs::copy(GetTestFile("com.android.apex.cts.shim.v2.apex"), data_dir.path);
+  ASSERT_TRUE(IsOk(instance.AddDataApex(data_dir.path)));
 
-  std::vector<std::string> dirs_to_scan{built_in_dir.path, data_dir.path};
-  auto all_apex = ScanAndGroupApexFiles(dirs_to_scan);
-  auto result = SelectApexForActivation(std::move(all_apex), instance);
+  auto all_apex = instance.AllApexFilesByName();
+  auto result = SelectApexForActivation(all_apex, instance);
   ASSERT_EQ(result.size(), 2u);
 
   auto apexd_test_file_v2 = ApexFile::Open(
@@ -143,7 +120,6 @@
   TemporaryDir built_in_dir;
   fs::copy(GetTestFile("apex.apexd_test.apex"), built_in_dir.path);
   fs::copy(GetTestFile("com.android.apex.cts.shim.apex"), built_in_dir.path);
-
   // Initialize pre-installed APEX information
   ApexFileRepository instance;
   ASSERT_TRUE(IsOk(instance.AddPreInstalledApex({built_in_dir.path})));
@@ -151,10 +127,11 @@
   TemporaryDir data_dir;
   fs::copy(GetTestFile("apex.apexd_test.apex"), data_dir.path);
   fs::copy(GetTestFile("com.android.apex.cts.shim.apex"), data_dir.path);
+  // Initialize ApexFile repo
+  ASSERT_TRUE(IsOk(instance.AddDataApex(data_dir.path)));
 
-  std::vector<std::string> dirs_to_scan{built_in_dir.path, data_dir.path};
-  auto all_apex = ScanAndGroupApexFiles(dirs_to_scan);
-  auto result = SelectApexForActivation(std::move(all_apex), instance);
+  auto all_apex = instance.AllApexFilesByName();
+  auto result = SelectApexForActivation(all_apex, instance);
   ASSERT_EQ(result.size(), 2u);
 
   auto apexd_test_file =
@@ -171,7 +148,6 @@
   fs::copy(
       GetTestFile("com.android.apex.test.sharedlibs_generated.v1.libvX.apex"),
       built_in_dir.path);
-
   // Initialize pre-installed APEX information
   ApexFileRepository instance;
   ASSERT_TRUE(IsOk(instance.AddPreInstalledApex({built_in_dir.path})));
@@ -180,10 +156,11 @@
   fs::copy(
       GetTestFile("com.android.apex.test.sharedlibs_generated.v2.libvY.apex"),
       data_dir.path);
+  // Initialize data APEX information
+  ASSERT_TRUE(IsOk(instance.AddDataApex(data_dir.path)));
 
-  std::vector<std::string> dirs_to_scan{built_in_dir.path, data_dir.path};
-  auto all_apex = ScanAndGroupApexFiles(dirs_to_scan);
-  auto result = SelectApexForActivation(std::move(all_apex), instance);
+  auto all_apex = instance.AllApexFilesByName();
+  auto result = SelectApexForActivation(all_apex, instance);
   ASSERT_EQ(result.size(), 2u);
 
   auto shared_lib_v1 = ApexFile::Open(StringPrintf(
@@ -204,11 +181,10 @@
       "%s/com.android.apex.compressed.v1.capex", built_in_dir.path));
 
   TemporaryDir decompression_dir, active_apex_dir;
-  std::vector<ApexFile> compressed_apex_list;
-  compressed_apex_list.emplace_back(std::move(*compressed_apex));
-  auto return_value =
-      ProcessCompressedApex(std::move(compressed_apex_list),
-                            decompression_dir.path, active_apex_dir.path);
+  std::vector<std::reference_wrapper<const ApexFile>> compressed_apex_list;
+  compressed_apex_list.emplace_back(std::cref(*compressed_apex));
+  auto return_value = ProcessCompressedApex(
+      compressed_apex_list, decompression_dir.path, active_apex_dir.path);
 
   std::string decompressed_file_path = StringPrintf(
       "%s/com.android.apex.compressed@1.apex", decompression_dir.path);
@@ -251,11 +227,10 @@
       built_in_dir.path));
 
   TemporaryDir decompression_dir, active_apex_dir;
-  std::vector<ApexFile> compressed_apex_list;
-  compressed_apex_list.emplace_back(std::move(*compressed_apex_mismatch_key));
-  auto return_value =
-      ProcessCompressedApex(std::move(compressed_apex_list),
-                            decompression_dir.path, active_apex_dir.path);
+  std::vector<std::reference_wrapper<const ApexFile>> compressed_apex_list;
+  compressed_apex_list.emplace_back(std::cref(*compressed_apex_mismatch_key));
+  auto return_value = ProcessCompressedApex(
+      compressed_apex_list, decompression_dir.path, active_apex_dir.path);
   ASSERT_EQ(return_value.size(), 0u);
 }