qt: fix romfs dumping for multiprogram applications

2 files changed, 65 insertions, 51 deletions

diff --git a/src/yuzu/main.cpp b/src/yuzu/main.cpp
index 33c9fd0af..81dfbca22 100644
--- a/src/yuzu/main.cpp
+++ b/src/yuzu/main.cpp
@@ -2573,50 +2573,41 @@ void GMainWindow::OnGameListDumpRomFS(u64 program_id, const std::string& game_pa
         return;
     }
 
-    FileSys::VirtualFile base_romfs;
-    if (loader->ReadRomFS(base_romfs) != Loader::ResultStatus::Success) {
-        failed();
-        return;
-    }
+    FileSys::VirtualFile packed_update_raw{};
+    loader->ReadUpdateRaw(packed_update_raw);
 
     const auto& installed = system->GetContentProvider();
-    const auto romfs_title_id = SelectRomFSDumpTarget(installed, program_id);
 
-    if (!romfs_title_id) {
+    u64 title_id{};
+    u8 raw_type{};
+    if (!SelectRomFSDumpTarget(installed, program_id, &title_id, &raw_type)) {
         failed();
         return;
     }
 
-    const auto type = *romfs_title_id == program_id ? FileSys::ContentRecordType::Program
-                                                    : FileSys::ContentRecordType::Data;
-    const auto base_nca = installed.GetEntry(*romfs_title_id, type);
+    const auto type = static_cast<FileSys::ContentRecordType>(raw_type);
+    const auto base_nca = installed.GetEntry(title_id, type);
     if (!base_nca) {
         failed();
         return;
     }
 
+    const auto base_romfs = base_nca->GetRomFS();
+    if (!base_romfs) {
+        failed();
+        return;
+    }
+
     const auto dump_dir =
         target == DumpRomFSTarget::Normal
             ? Common::FS::GetYuzuPath(Common::FS::YuzuPath::DumpDir)
             : Common::FS::GetYuzuPath(Common::FS::YuzuPath::SDMCDir) / "atmosphere" / "contents";
-    const auto romfs_dir = fmt::format("{:016X}/romfs", *romfs_title_id);
+    const auto romfs_dir = fmt::format("{:016X}/romfs", title_id);
 
     const auto path = Common::FS::PathToUTF8String(dump_dir / romfs_dir);
 
-    FileSys::VirtualFile romfs;
-
-    if (*romfs_title_id == program_id) {
-        const FileSys::PatchManager pm{program_id, system->GetFileSystemController(), installed};
-        romfs = pm.PatchRomFS(base_nca.get(), base_romfs, type, nullptr, false);
-    } else {
-        romfs = installed.GetEntry(*romfs_title_id, type)->GetRomFS();
-    }
-
-    const auto extracted = FileSys::ExtractRomFS(romfs, FileSys::RomFSExtractionType::Full);
-    if (extracted == nullptr) {
-        failed();
-        return;
-    }
+    const FileSys::PatchManager pm{title_id, system->GetFileSystemController(), installed};
+    auto romfs = pm.PatchRomFS(base_nca.get(), base_romfs, type, packed_update_raw, false);
 
     const auto out = VfsFilesystemCreateDirectoryWrapper(vfs, path, FileSys::Mode::ReadWrite);
 
@@ -2640,6 +2631,12 @@ void GMainWindow::OnGameListDumpRomFS(u64 program_id, const std::string& game_pa
         return;
     }
 
+    const auto extracted = FileSys::ExtractRomFS(romfs, FileSys::RomFSExtractionType::Full);
+    if (extracted == nullptr) {
+        failed();
+        return;
+    }
+
     const auto full = res == selections.constFirst();
     const auto entry_size = CalculateRomFSEntrySize(extracted, full);
 
@@ -4350,28 +4347,41 @@ bool GMainWindow::CheckSystemArchiveDecryption() {
     return mii_nca->GetRomFS().get() != nullptr;
 }
 
-std::optional<u64> GMainWindow::SelectRomFSDumpTarget(const FileSys::ContentProvider& installed,
-                                                      u64 program_id) {
-    const auto dlc_entries =
-        installed.ListEntriesFilter(FileSys::TitleType::AOC, FileSys::ContentRecordType::Data);
-    std::vector<FileSys::ContentProviderEntry> dlc_match;
-    dlc_match.reserve(dlc_entries.size());
-    std::copy_if(dlc_entries.begin(), dlc_entries.end(), std::back_inserter(dlc_match),
-                 [&program_id, &installed](const FileSys::ContentProviderEntry& entry) {
-                     return FileSys::GetBaseTitleID(entry.title_id) == program_id &&
-                            installed.GetEntry(entry)->GetStatus() == Loader::ResultStatus::Success;
-                 });
-
-    std::vector<u64> romfs_tids;
-    romfs_tids.push_back(program_id);
-    for (const auto& entry : dlc_match) {
-        romfs_tids.push_back(entry.title_id);
-    }
-
-    if (romfs_tids.size() > 1) {
-        QStringList list{QStringLiteral("Base")};
-        for (std::size_t i = 1; i < romfs_tids.size(); ++i) {
-            list.push_back(QStringLiteral("DLC %1").arg(romfs_tids[i] & 0x7FF));
+bool GMainWindow::SelectRomFSDumpTarget(const FileSys::ContentProvider& installed, u64 program_id,
+                                        u64* selected_title_id, u8* selected_content_record_type) {
+    using ContentInfo = std::pair<FileSys::TitleType, FileSys::ContentRecordType>;
+    boost::container::flat_map<u64, ContentInfo> available_title_ids;
+
+    const auto RetrieveEntries = [&](FileSys::TitleType title_type,
+                                     FileSys::ContentRecordType record_type) {
+        const auto entries = installed.ListEntriesFilter(title_type, record_type);
+        for (const auto& entry : entries) {
+            if (FileSys::GetBaseTitleID(entry.title_id) == program_id &&
+                installed.GetEntry(entry)->GetStatus() == Loader::ResultStatus::Success) {
+                available_title_ids[entry.title_id] = {title_type, record_type};
+            }
+        }
+    };
+
+    RetrieveEntries(FileSys::TitleType::Application, FileSys::ContentRecordType::Program);
+    RetrieveEntries(FileSys::TitleType::AOC, FileSys::ContentRecordType::Data);
+
+    if (available_title_ids.empty()) {
+        return false;
+    }
+
+    size_t title_index = 0;
+
+    if (available_title_ids.size() > 1) {
+        QStringList list;
+        for (auto& [title_id, content_info] : available_title_ids) {
+            const auto hex_title_id = QString::fromStdString(fmt::format("{:X}", title_id));
+            if (content_info.first == FileSys::TitleType::Application) {
+                list.push_back(QStringLiteral("Application [%1]").arg(hex_title_id));
+            } else {
+                list.push_back(
+                    QStringLiteral("DLC %1 [%2]").arg(title_id & 0x7FF).arg(hex_title_id));
+            }
         }
 
         bool ok;
@@ -4379,13 +4389,16 @@ std::optional<u64> GMainWindow::SelectRomFSDumpTarget(const FileSys::ContentProv
             this, tr("Select RomFS Dump Target"),
             tr("Please select which RomFS you would like to dump."), list, 0, false, &ok);
         if (!ok) {
-            return {};
+            return false;
         }
 
-        return romfs_tids[list.indexOf(res)];
+        title_index = list.indexOf(res);
     }
 
-    return program_id;
+    const auto selected_info = available_title_ids.nth(title_index);
+    *selected_title_id = selected_info->first;
+    *selected_content_record_type = static_cast<u8>(selected_info->second.second);
+    return true;
 }
 
 bool GMainWindow::ConfirmClose() {
diff --git a/src/yuzu/main.h b/src/yuzu/main.h
index 1b7055122..668dbc3b1 100644
--- a/src/yuzu/main.h
+++ b/src/yuzu/main.h
@@ -375,7 +375,8 @@ private:
     void RemoveAllTransferableShaderCaches(u64 program_id);
     void RemoveCustomConfiguration(u64 program_id, const std::string& game_path);
     void RemoveCacheStorage(u64 program_id);
-    std::optional<u64> SelectRomFSDumpTarget(const FileSys::ContentProvider&, u64 program_id);
+    bool SelectRomFSDumpTarget(const FileSys::ContentProvider&, u64 program_id,
+                               u64* selected_title_id, u8* selected_content_record_type);
     InstallResult InstallNSPXCI(const QString& filename);
     InstallResult InstallNCA(const QString& filename);
     void MigrateConfigFiles();