core: Support multiple modules per patcher

3 files changed, 80 insertions, 39 deletions

diff --git a/src/core/loader/deconstructed_rom_directory.cpp b/src/core/loader/deconstructed_rom_directory.cpp
index c9f8707b7..b2173f697 100644
--- a/src/core/loader/deconstructed_rom_directory.cpp
+++ b/src/core/loader/deconstructed_rom_directory.cpp
@@ -19,8 +19,54 @@
 #include "core/arm/nce/patcher.h"
 #endif
 
+#ifndef HAS_NCE
+namespace Core::NCE {
+class Patcher {};
+} // namespace Core::NCE
+#endif
+
 namespace Loader {
 
+struct PatchCollection {
+    explicit PatchCollection(bool is_application_) : is_application{is_application_} {
+        module_patcher_indices.fill(-1);
+        patchers.emplace_back();
+    }
+
+    std::vector<Core::NCE::Patcher>* GetPatchers() {
+        if (is_application && Settings::IsNceEnabled()) {
+            return &patchers;
+        }
+        return nullptr;
+    }
+
+    size_t GetTotalPatchSize() const {
+        size_t total_size{};
+#ifdef HAS_NCE
+        for (auto& patcher : patchers) {
+            total_size += patcher.GetSectionSize();
+        }
+#endif
+        return total_size;
+    }
+
+    void SaveIndex(size_t module) {
+        module_patcher_indices[module] = static_cast<s32>(patchers.size() - 1);
+    }
+
+    s32 GetIndex(size_t module) const {
+        return module_patcher_indices[module];
+    }
+
+    s32 GetLastIndex() const {
+        return static_cast<s32>(patchers.size()) - 1;
+    }
+
+    bool is_application;
+    std::vector<Core::NCE::Patcher> patchers;
+    std::array<s32, 13> module_patcher_indices{};
+};
+
 AppLoader_DeconstructedRomDirectory::AppLoader_DeconstructedRomDirectory(FileSys::VirtualFile file_,
                                                                          bool override_update_)
     : AppLoader(std::move(file_)), override_update(override_update_), is_hbl(false) {
@@ -142,18 +188,7 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
     std::size_t code_size{};
 
     // Define an nce patch context for each potential module.
-#ifdef HAS_NCE
-    std::array<Core::NCE::Patcher, 13> module_patchers;
-#endif
-
-    const auto GetPatcher = [&](size_t i) -> Core::NCE::Patcher* {
-#ifdef HAS_NCE
-        if (is_application && Settings::IsNceEnabled()) {
-            return &module_patchers[i];
-        }
-#endif
-        return nullptr;
-    };
+    PatchCollection patch_ctx{is_application};
 
     // Use the NSO module loader to figure out the code layout
     for (size_t i = 0; i < static_modules.size(); i++) {
@@ -164,13 +199,14 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
         }
 
         const bool should_pass_arguments = std::strcmp(module, "rtld") == 0;
-        const auto tentative_next_load_addr =
-            AppLoader_NSO::LoadModule(process, system, *module_file, code_size,
-                                      should_pass_arguments, false, {}, GetPatcher(i));
+        const auto tentative_next_load_addr = AppLoader_NSO::LoadModule(
+            process, system, *module_file, code_size, should_pass_arguments, false, {},
+            patch_ctx.GetPatchers(), patch_ctx.GetLastIndex());
         if (!tentative_next_load_addr) {
             return {ResultStatus::ErrorLoadingNSO, {}};
         }
 
+        patch_ctx.SaveIndex(i);
         code_size = *tentative_next_load_addr;
     }
 
@@ -184,6 +220,9 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
         return 0;
     }();
 
+    // Add patch size to the total module size
+    code_size += patch_ctx.GetTotalPatchSize();
+
     // Setup the process code layout
     if (process.LoadFromMetadata(metadata, code_size, fastmem_base, is_hbl).IsError()) {
         return {ResultStatus::ErrorUnableToParseKernelMetadata, {}};
@@ -204,9 +243,9 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
 
         const VAddr load_addr{next_load_addr};
         const bool should_pass_arguments = std::strcmp(module, "rtld") == 0;
-        const auto tentative_next_load_addr =
-            AppLoader_NSO::LoadModule(process, system, *module_file, load_addr,
-                                      should_pass_arguments, true, pm, GetPatcher(i));
+        const auto tentative_next_load_addr = AppLoader_NSO::LoadModule(
+            process, system, *module_file, load_addr, should_pass_arguments, true, pm,
+            patch_ctx.GetPatchers(), patch_ctx.GetIndex(i));
         if (!tentative_next_load_addr) {
             return {ResultStatus::ErrorLoadingNSO, {}};
         }
diff --git a/src/core/loader/nso.cpp b/src/core/loader/nso.cpp
index b053a0d14..583b7e927 100644
--- a/src/core/loader/nso.cpp
+++ b/src/core/loader/nso.cpp
@@ -77,7 +77,8 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
                                                const FileSys::VfsFile& nso_file, VAddr load_base,
                                                bool should_pass_arguments, bool load_into_process,
                                                std::optional<FileSys::PatchManager> pm,
-                                               Core::NCE::Patcher* patch) {
+                                               std::vector<Core::NCE::Patcher>* patches,
+                                               s32 patch_index) {
     if (nso_file.GetSize() < sizeof(NSOHeader)) {
         return std::nullopt;
     }
@@ -94,8 +95,11 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
     // Allocate some space at the beginning if we are patching in PreText mode.
     const size_t module_start = [&]() -> size_t {
 #ifdef HAS_NCE
-        if (patch && patch->GetPatchMode() == Core::NCE::PatchMode::PreText) {
-            return patch->GetSectionSize();
+        if (patches && load_into_process) {
+            auto* patch = &patches->operator[](patch_index);
+            if (patch->GetPatchMode() == Core::NCE::PatchMode::PreText) {
+                return patch->GetSectionSize();
+            }
         }
 #endif
         return 0;
@@ -160,27 +164,24 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
 #ifdef HAS_NCE
     // If we are computing the process code layout and using nce backend, patch.
     const auto& code = codeset.CodeSegment();
-    if (patch && patch->GetPatchMode() == Core::NCE::PatchMode::None) {
+    auto* patch = patches ? &patches->operator[](patch_index) : nullptr;
+    if (patch && !load_into_process) {
         // Patch SVCs and MRS calls in the guest code
-        patch->PatchText(program_image, code);
-
-        // Add patch section size to the module size.
-        image_size += static_cast<u32>(patch->GetSectionSize());
+        while (!patch->PatchText(program_image, code)) {
+            patch = &patches->emplace_back();
+        }
     } else if (patch) {
         // Relocate code patch and copy to the program_image.
-        patch->RelocateAndCopy(load_base, code, program_image, &process.GetPostHandlers());
-
-        // Update patch section.
-        auto& patch_segment = codeset.PatchSegment();
-        patch_segment.addr =
-            patch->GetPatchMode() == Core::NCE::PatchMode::PreText ? 0 : image_size;
-        patch_segment.size = static_cast<u32>(patch->GetSectionSize());
-
-        // Add patch section size to the module size. In PreText mode image_size
-        // already contains the patch segment as part of module_start.
-        if (patch->GetPatchMode() == Core::NCE::PatchMode::PostData) {
-            image_size += patch_segment.size;
+        if (patch->RelocateAndCopy(load_base, code, program_image, &process.GetPostHandlers())) {
+            // Update patch section.
+            auto& patch_segment = codeset.PatchSegment();
+            patch_segment.addr =
+                patch->GetPatchMode() == Core::NCE::PatchMode::PreText ? 0 : image_size;
+            patch_segment.size = static_cast<u32>(patch->GetSectionSize());
         }
+
+        // Refresh image_size to take account the patch section if it was added by RelocateAndCopy
+        image_size = static_cast<u32>(program_image.size());
     }
 #endif
 
diff --git a/src/core/loader/nso.h b/src/core/loader/nso.h
index 29b86ed4c..6356697e3 100644
--- a/src/core/loader/nso.h
+++ b/src/core/loader/nso.h
@@ -93,7 +93,8 @@ public:
                                            const FileSys::VfsFile& nso_file, VAddr load_base,
                                            bool should_pass_arguments, bool load_into_process,
                                            std::optional<FileSys::PatchManager> pm = {},
-                                           Core::NCE::Patcher* patch = nullptr);
+                                           std::vector<Core::NCE::Patcher>* patches = nullptr,
+                                           s32 patch_index = -1);
 
     LoadResult Load(Kernel::KProcess& process, Core::System& system) override;