diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index dbae70e..f535557 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -1426,25 +1426,15 @@
     if (profile_compilation_info_ != nullptr && IsAppImage()) {
       Runtime* runtime = Runtime::Current();
       CHECK(runtime != nullptr);
-      std::set<DexCacheResolvedClasses> resolved_classes(
-          profile_compilation_info_->GetResolvedClasses());
-
       // Filter out class path classes since we don't want to include these in the image.
       std::unordered_set<std::string> dex_files_locations;
       for (const DexFile* dex_file : dex_files_) {
         dex_files_locations.insert(dex_file->GetLocation());
       }
-      for (auto it = resolved_classes.begin(); it != resolved_classes.end(); ) {
-        if (dex_files_locations.find(it->GetDexLocation()) == dex_files_locations.end()) {
-          VLOG(compiler) << "Removed profile samples for non-app dex file " << it->GetDexLocation();
-          it = resolved_classes.erase(it);
-        } else {
-          ++it;
-        }
-      }
-
+      std::set<DexCacheResolvedClasses> resolved_classes(
+          profile_compilation_info_->GetResolvedClasses(dex_files_locations));
       image_classes_.reset(new std::unordered_set<std::string>(
-          runtime->GetClassLinker()->GetClassDescriptorsForProfileKeys(resolved_classes)));
+          runtime->GetClassLinker()->GetClassDescriptorsForResolvedClasses(resolved_classes)));
       VLOG(compiler) << "Loaded " << image_classes_->size()
                      << " image class descriptors from profile";
       if (VLOG_IS_ON(compiler)) {
diff --git a/dex2oat/dex2oat_test.cc b/dex2oat/dex2oat_test.cc
index b79050e..e7277bc 100644
--- a/dex2oat/dex2oat_test.cc
+++ b/dex2oat/dex2oat_test.cc
@@ -63,6 +63,7 @@
       oat_file.reset(OS::CreateEmptyFile(odex_location.c_str()));
       CHECK(oat_file != nullptr) << odex_location;
       args.push_back("--oat-fd=" + std::to_string(oat_file->Fd()));
+      args.push_back("--oat-location=" + odex_location);
     } else {
       args.push_back("--oat-file=" + odex_location);
     }
@@ -583,13 +584,16 @@
   // Emits a profile with a single dex file with the given location and a single class index of 1.
   void GenerateProfile(const std::string& test_profile,
                        const std::string& dex_location,
+                       size_t num_classes,
                        uint32_t checksum) {
     int profile_test_fd = open(test_profile.c_str(), O_CREAT | O_TRUNC | O_WRONLY, 0644);
     CHECK_GE(profile_test_fd, 0);
 
     ProfileCompilationInfo info;
     std::string profile_key = ProfileCompilationInfo::GetProfileDexFileKey(dex_location);
-    info.AddClassIndex(profile_key, checksum, dex::TypeIndex(1));
+    for (size_t i = 0; i < num_classes; ++i) {
+      info.AddClassIndex(profile_key, checksum, dex::TypeIndex(1 + i));
+    }
     bool result = info.Save(profile_test_fd);
     close(profile_test_fd);
     ASSERT_TRUE(result);
@@ -597,7 +601,9 @@
 
   void CompileProfileOdex(const std::string& dex_location,
                           const std::string& odex_location,
+                          const std::string& app_image_file_name,
                           bool use_fd,
+                          size_t num_profile_classes,
                           const std::vector<std::string>& extra_args = {}) {
     const std::string profile_location = GetScratchDir() + "/primary.prof";
     const char* location = dex_location.c_str();
@@ -606,33 +612,86 @@
     ASSERT_TRUE(DexFile::Open(location, location, true, &error_msg, &dex_files));
     EXPECT_EQ(dex_files.size(), 1U);
     std::unique_ptr<const DexFile>& dex_file = dex_files[0];
-    GenerateProfile(profile_location, dex_location, dex_file->GetLocationChecksum());
+    GenerateProfile(profile_location,
+                    dex_location,
+                    num_profile_classes,
+                    dex_file->GetLocationChecksum());
     std::vector<std::string> copy(extra_args);
     copy.push_back("--profile-file=" + profile_location);
+    std::unique_ptr<File> app_image_file;
+    if (!app_image_file_name.empty()) {
+      if (use_fd) {
+        app_image_file.reset(OS::CreateEmptyFile(app_image_file_name.c_str()));
+        copy.push_back("--app-image-fd=" + std::to_string(app_image_file->Fd()));
+      } else {
+        copy.push_back("--app-image-file=" + app_image_file_name);
+      }
+    }
     GenerateOdexForTest(dex_location,
                         odex_location,
                         CompilerFilter::kSpeedProfile,
                         copy,
                         /* expect_success */ true,
                         use_fd);
+    if (app_image_file != nullptr) {
+      ASSERT_EQ(app_image_file->FlushCloseOrErase(), 0) << "Could not flush and close art file";
+    }
   }
 
-  void RunTest() {
+  uint64_t GetImageSize(const std::string& image_file_name) {
+    EXPECT_FALSE(image_file_name.empty());
+    std::unique_ptr<File> file(OS::OpenFileForReading(image_file_name.c_str()));
+    CHECK(file != nullptr);
+    ImageHeader image_header;
+    const bool success = file->ReadFully(&image_header, sizeof(image_header));
+    CHECK(success);
+    CHECK(image_header.IsValid());
+    ReaderMutexLock mu(Thread::Current(), *Locks::mutator_lock_);
+    return image_header.GetImageSize();
+  }
+
+  void RunTest(bool app_image) {
     std::string dex_location = GetScratchDir() + "/DexNoOat.jar";
     std::string odex_location = GetOdexDir() + "/DexOdexNoOat.odex";
+    std::string app_image_file = app_image ? (GetOdexDir() + "/DexOdexNoOat.art"): "";
     Copy(GetDexSrc2(), dex_location);
 
-    CompileProfileOdex(dex_location, odex_location, /* use_fd */ false);
+    uint64_t image_file_empty_profile = 0;
+    if (app_image) {
+      CompileProfileOdex(dex_location,
+                         odex_location,
+                         app_image_file,
+                         /* use_fd */ false,
+                         /* num_profile_classes */ 0);
+      CheckValidity();
+      ASSERT_TRUE(success_);
+      // Don't check the result since CheckResult relies on the class being in the profile.
+      image_file_empty_profile = GetImageSize(app_image_file);
+      EXPECT_GT(image_file_empty_profile, 0u);
+    }
 
+    // Small profile.
+    CompileProfileOdex(dex_location,
+                       odex_location,
+                       app_image_file,
+                       /* use_fd */ false,
+                       /* num_profile_classes */ 1);
     CheckValidity();
     ASSERT_TRUE(success_);
-    CheckResult(dex_location, odex_location);
+    CheckResult(dex_location, odex_location, app_image_file);
+
+    if (app_image) {
+      // Test that the profile made a difference by adding more classes.
+      const uint64_t image_file_small_profile = GetImageSize(app_image_file);
+      CHECK_LT(image_file_empty_profile, image_file_small_profile);
+    }
   }
 
   void RunTestVDex() {
     std::string dex_location = GetScratchDir() + "/DexNoOat.jar";
     std::string odex_location = GetOdexDir() + "/DexOdexNoOat.odex";
     std::string vdex_location = GetOdexDir() + "/DexOdexNoOat.vdex";
+    std::string app_image_file_name = GetOdexDir() + "/DexOdexNoOat.art";
     Copy(GetDexSrc2(), dex_location);
 
     std::unique_ptr<File> vdex_file1(OS::CreateEmptyFile(vdex_location.c_str()));
@@ -643,7 +702,9 @@
       std::string output_vdex = StringPrintf("--output-vdex-fd=%d", vdex_file1->Fd());
       CompileProfileOdex(dex_location,
                          odex_location,
+                         app_image_file_name,
                          /* use_fd */ true,
+                         /* num_profile_classes */ 1,
                          { input_vdex, output_vdex });
       EXPECT_GT(vdex_file1->GetLength(), 0u);
     }
@@ -652,17 +713,21 @@
       std::string output_vdex = StringPrintf("--output-vdex-fd=%d", vdex_file2.GetFd());
       CompileProfileOdex(dex_location,
                          odex_location,
+                         app_image_file_name,
                          /* use_fd */ true,
+                         /* num_profile_classes */ 1,
                          { input_vdex, output_vdex });
       EXPECT_GT(vdex_file2.GetFile()->GetLength(), 0u);
     }
     ASSERT_EQ(vdex_file1->FlushCloseOrErase(), 0) << "Could not flush and close vdex file";
     CheckValidity();
     ASSERT_TRUE(success_);
-    CheckResult(dex_location, odex_location);
+    CheckResult(dex_location, odex_location, app_image_file_name);
   }
 
-  void CheckResult(const std::string& dex_location, const std::string& odex_location) {
+  void CheckResult(const std::string& dex_location,
+                   const std::string& odex_location,
+                   const std::string& app_image_file_name) {
     // Host/target independent checks.
     std::string error_msg;
     std::unique_ptr<OatFile> odex_file(OatFile::Open(odex_location.c_str(),
@@ -698,6 +763,16 @@
     }
 
     EXPECT_EQ(odex_file->GetCompilerFilter(), CompilerFilter::kSpeedProfile);
+
+    if (!app_image_file_name.empty()) {
+      // Go peek at the image header to make sure it was large enough to contain the class.
+      std::unique_ptr<File> file(OS::OpenFileForReading(app_image_file_name.c_str()));
+      ImageHeader image_header;
+      bool success = file->ReadFully(&image_header, sizeof(image_header));
+      ASSERT_TRUE(success);
+      ASSERT_TRUE(image_header.IsValid());
+      EXPECT_GT(image_header.GetImageSection(ImageHeader::kSectionObjects).Size(), 0u);
+    }
   }
 
   // Check whether the dex2oat run was really successful.
@@ -720,7 +795,11 @@
 };
 
 TEST_F(Dex2oatLayoutTest, TestLayout) {
-  RunTest();
+  RunTest(/* app-image */ false);
+}
+
+TEST_F(Dex2oatLayoutTest, TestLayoutAppImage) {
+  RunTest(/* app-image */ true);
 }
 
 TEST_F(Dex2oatLayoutTest, TestVdexLayout) {
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index cac5449..d232714 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -8933,7 +8933,7 @@
   return ret;
 }
 
-std::unordered_set<std::string> ClassLinker::GetClassDescriptorsForProfileKeys(
+std::unordered_set<std::string> ClassLinker::GetClassDescriptorsForResolvedClasses(
     const std::set<DexCacheResolvedClasses>& classes) {
   ScopedTrace trace(__PRETTY_FUNCTION__);
   std::unordered_set<std::string> ret;
@@ -8948,14 +8948,13 @@
       if (dex_cache != nullptr) {
         const DexFile* dex_file = dex_cache->GetDexFile();
         // There could be duplicates if two dex files with the same location are mapped.
-        location_to_dex_file.emplace(
-            ProfileCompilationInfo::GetProfileDexFileKey(dex_file->GetLocation()), dex_file);
+        location_to_dex_file.emplace(dex_file->GetLocation(), dex_file);
       }
     }
   }
   for (const DexCacheResolvedClasses& info : classes) {
-    const std::string& profile_key = info.GetDexLocation();
-    auto found = location_to_dex_file.find(profile_key);
+    const std::string& location = info.GetDexLocation();
+    auto found = location_to_dex_file.find(location);
     if (found != location_to_dex_file.end()) {
       const DexFile* dex_file = found->second;
       VLOG(profiler) << "Found opened dex file for " << dex_file->GetLocation() << " with "
@@ -8967,7 +8966,7 @@
         ret.insert(descriptor);
       }
     } else {
-      VLOG(class_linker) << "Failed to find opened dex file for profile key " << profile_key;
+      VLOG(class_linker) << "Failed to find opened dex file for location " << location;
     }
   }
   return ret;
diff --git a/runtime/class_linker.h b/runtime/class_linker.h
index 33eed3c..a5d26c7 100644
--- a/runtime/class_linker.h
+++ b/runtime/class_linker.h
@@ -617,7 +617,8 @@
   std::set<DexCacheResolvedClasses> GetResolvedClasses(bool ignore_boot_classes)
       REQUIRES(!Locks::dex_lock_);
 
-  std::unordered_set<std::string> GetClassDescriptorsForProfileKeys(
+  // Returns the class descriptors for loaded dex files.
+  std::unordered_set<std::string> GetClassDescriptorsForResolvedClasses(
       const std::set<DexCacheResolvedClasses>& classes)
       REQUIRES(!Locks::dex_lock_);
 
diff --git a/runtime/jit/profile_compilation_info.cc b/runtime/jit/profile_compilation_info.cc
index 627cc93..3467b60 100644
--- a/runtime/jit/profile_compilation_info.cc
+++ b/runtime/jit/profile_compilation_info.cc
@@ -1039,15 +1039,22 @@
   return info_.Equals(other.info_);
 }
 
-std::set<DexCacheResolvedClasses> ProfileCompilationInfo::GetResolvedClasses() const {
+std::set<DexCacheResolvedClasses> ProfileCompilationInfo::GetResolvedClasses(
+    const std::unordered_set<std::string>& dex_files_locations) const {
+  std::unordered_map<std::string, std::string> key_to_location_map;
+  for (const std::string& location : dex_files_locations) {
+    key_to_location_map.emplace(GetProfileDexFileKey(location), location);
+  }
   std::set<DexCacheResolvedClasses> ret;
   for (auto&& pair : info_) {
     const std::string& profile_key = pair.first;
-    const DexFileData& data = pair.second;
-    // TODO: Is it OK to use the same location for both base and dex location here?
-    DexCacheResolvedClasses classes(profile_key, profile_key, data.checksum);
-    classes.AddClasses(data.class_set.begin(), data.class_set.end());
-    ret.insert(classes);
+    auto it = key_to_location_map.find(profile_key);
+    if (it != key_to_location_map.end()) {
+      const DexFileData& data = pair.second;
+      DexCacheResolvedClasses classes(it->second, it->second, data.checksum);
+      classes.AddClasses(data.class_set.begin(), data.class_set.end());
+      ret.insert(classes);
+    }
   }
   return ret;
 }
diff --git a/runtime/jit/profile_compilation_info.h b/runtime/jit/profile_compilation_info.h
index 4bfbfcd..cf556bf 100644
--- a/runtime/jit/profile_compilation_info.h
+++ b/runtime/jit/profile_compilation_info.h
@@ -218,9 +218,8 @@
   bool Equals(const ProfileCompilationInfo& other);
 
   // Return the class descriptors for all of the classes in the profiles' class sets.
-  // Note the dex location is actually the profile key, the caller needs to call back in to the
-  // profile info stuff to generate a map back to the dex location.
-  std::set<DexCacheResolvedClasses> GetResolvedClasses() const;
+  std::set<DexCacheResolvedClasses> GetResolvedClasses(
+      const std::unordered_set<std::string>& dex_files_locations) const;
 
   // Clear the resolved classes from the current object.
   void ClearResolvedClasses();