5 files changed, 162 insertions, 26 deletions

diff --git a/src/core/loader/deconstructed_rom_directory.cpp b/src/core/loader/deconstructed_rom_directory.cpp
index 48c0edaea..e7fc8f438 100644
--- a/src/core/loader/deconstructed_rom_directory.cpp
+++ b/src/core/loader/deconstructed_rom_directory.cpp
@@ -3,6 +3,7 @@
 
 #include <cstring>
 #include "common/logging/log.h"
+#include "common/settings.h"
 #include "core/core.h"
 #include "core/file_sys/content_archive.h"
 #include "core/file_sys/control_metadata.h"
@@ -14,6 +15,10 @@
 #include "core/loader/deconstructed_rom_directory.h"
 #include "core/loader/nso.h"
 
+#ifdef ARCHITECTURE_arm64
+#include "core/arm/nce/patch.h"
+#endif
+
 namespace Loader {
 
 AppLoader_DeconstructedRomDirectory::AppLoader_DeconstructedRomDirectory(FileSys::VirtualFile file_,
@@ -124,21 +129,41 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
     }
     metadata.Print();
 
-    const auto static_modules = {"rtld",    "main",    "subsdk0", "subsdk1", "subsdk2",
-                                 "subsdk3", "subsdk4", "subsdk5", "subsdk6", "subsdk7",
-                                 "subsdk8", "subsdk9", "sdk"};
+    // Enable NCE only for 64-bit programs.
+    Settings::SetNceEnabled(metadata.Is64BitProgram());
+
+    const std::array static_modules = {"rtld",    "main",    "subsdk0", "subsdk1", "subsdk2",
+                                       "subsdk3", "subsdk4", "subsdk5", "subsdk6", "subsdk7",
+                                       "subsdk8", "subsdk9", "sdk"};
 
-    // Use the NSO module loader to figure out the code layout
     std::size_t code_size{};
-    for (const auto& module : static_modules) {
+
+    // Define an nce patch context for each potential module.
+#ifdef ARCHITECTURE_arm64
+    std::array<Core::NCE::Patcher, 13> module_patchers;
+#endif
+
+    const auto GetPatcher = [&](size_t i) -> Core::NCE::Patcher* {
+#ifdef ARCHITECTURE_arm64
+        if (Settings::IsNceEnabled()) {
+            return &module_patchers[i];
+        }
+#endif
+        return nullptr;
+    };
+
+    // Use the NSO module loader to figure out the code layout
+    for (size_t i = 0; i < static_modules.size(); i++) {
+        const auto& module = static_modules[i];
         const FileSys::VirtualFile module_file{dir->GetFile(module)};
         if (!module_file) {
             continue;
         }
 
         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);
+        const auto tentative_next_load_addr =
+            AppLoader_NSO::LoadModule(process, system, *module_file, code_size,
+                                      should_pass_arguments, false, {}, GetPatcher(i));
         if (!tentative_next_load_addr) {
             return {ResultStatus::ErrorLoadingNSO, {}};
         }
@@ -146,8 +171,18 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
         code_size = *tentative_next_load_addr;
     }
 
+    // Enable direct memory mapping in case of NCE.
+    const u64 fastmem_base = [&]() -> size_t {
+        if (Settings::IsNceEnabled()) {
+            auto& buffer = system.DeviceMemory().buffer;
+            buffer.EnableDirectMappedAddress();
+            return reinterpret_cast<u64>(buffer.VirtualBasePointer());
+        }
+        return 0;
+    }();
+
     // Setup the process code layout
-    if (process.LoadFromMetadata(metadata, code_size, 0, is_hbl).IsError()) {
+    if (process.LoadFromMetadata(metadata, code_size, fastmem_base, is_hbl).IsError()) {
         return {ResultStatus::ErrorUnableToParseKernelMetadata, {}};
     }
 
@@ -157,7 +192,8 @@ AppLoader_DeconstructedRomDirectory::LoadResult AppLoader_DeconstructedRomDirect
     VAddr next_load_addr{base_address};
     const FileSys::PatchManager pm{metadata.GetTitleID(), system.GetFileSystemController(),
                                    system.GetContentProvider()};
-    for (const auto& module : static_modules) {
+    for (size_t i = 0; i < static_modules.size(); i++) {
+        const auto& module = static_modules[i];
         const FileSys::VirtualFile module_file{dir->GetFile(module)};
         if (!module_file) {
             continue;
@@ -165,15 +201,16 @@ 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);
+        const auto tentative_next_load_addr =
+            AppLoader_NSO::LoadModule(process, system, *module_file, load_addr,
+                                      should_pass_arguments, true, pm, GetPatcher(i));
         if (!tentative_next_load_addr) {
             return {ResultStatus::ErrorLoadingNSO, {}};
         }
 
         next_load_addr = *tentative_next_load_addr;
         modules.insert_or_assign(load_addr, module);
-        LOG_DEBUG(Loader, "loaded module {} @ 0x{:X}", module, load_addr);
+        LOG_DEBUG(Loader, "loaded module {} @ {:#X}", module, load_addr);
     }
 
     // Find the RomFS by searching for a ".romfs" file in this directory
diff --git a/src/core/loader/nro.cpp b/src/core/loader/nro.cpp
index dfed296a5..49cf90317 100644
--- a/src/core/loader/nro.cpp
+++ b/src/core/loader/nro.cpp
@@ -22,6 +22,10 @@
 #include "core/loader/nso.h"
 #include "core/memory.h"
 
+#ifdef ARCHITECTURE_arm64
+#include "core/arm/nce/patch.h"
+#endif
+
 namespace Loader {
 
 struct NroSegmentHeader {
@@ -139,7 +143,8 @@ static constexpr u32 PageAlignSize(u32 size) {
     return static_cast<u32>((size + Core::Memory::YUZU_PAGEMASK) & ~Core::Memory::YUZU_PAGEMASK);
 }
 
-static bool LoadNroImpl(Kernel::KProcess& process, const std::vector<u8>& data) {
+static bool LoadNroImpl(Core::System& system, Kernel::KProcess& process,
+                        const std::vector<u8>& data) {
     if (data.size() < sizeof(NroHeader)) {
         return {};
     }
@@ -195,14 +200,60 @@ static bool LoadNroImpl(Kernel::KProcess& process, const std::vector<u8>& data)
     codeset.DataSegment().size += bss_size;
     program_image.resize(static_cast<u32>(program_image.size()) + bss_size);
 
+#ifdef ARCHITECTURE_arm64
+    const auto& code = codeset.CodeSegment();
+
+    // NROs are always 64-bit programs.
+    Settings::SetNceEnabled(true);
+
+    // Create NCE patcher
+    Core::NCE::Patcher patch{};
+    size_t image_size = program_image.size();
+
+    if (Settings::IsNceEnabled()) {
+        // Patch SVCs and MRS calls in the guest code
+        patch.PatchText(program_image, code);
+
+        // We only support PostData patching for NROs.
+        ASSERT(patch.Mode() == Core::NCE::PatchMode::PostData);
+
+        // Update patch section.
+        auto& patch_segment = codeset.PatchSegment();
+        patch_segment.addr = image_size;
+        patch_segment.size = static_cast<u32>(patch.SectionSize());
+
+        // Add patch section size to the module size.
+        image_size += patch_segment.size;
+    }
+#endif
+
+    // Enable direct memory mapping in case of NCE.
+    const u64 fastmem_base = [&]() -> size_t {
+        if (Settings::IsNceEnabled()) {
+            auto& buffer = system.DeviceMemory().buffer;
+            buffer.EnableDirectMappedAddress();
+            return reinterpret_cast<u64>(buffer.VirtualBasePointer());
+        }
+        return 0;
+    }();
+
     // Setup the process code layout
     if (process
-            .LoadFromMetadata(FileSys::ProgramMetadata::GetDefault(), program_image.size(), 0,
+            .LoadFromMetadata(FileSys::ProgramMetadata::GetDefault(), image_size, fastmem_base,
                               false)
             .IsError()) {
         return false;
     }
 
+    // Relocate code patch and copy to the program_image if running under NCE.
+    // This needs to be after LoadFromMetadata so we can use the process entry point.
+#ifdef ARCHITECTURE_arm64
+    if (Settings::IsNceEnabled()) {
+        patch.RelocateAndCopy(process.GetEntryPoint(), code, program_image,
+                              &process.GetPostHandlers());
+    }
+#endif
+
     // Load codeset for current process
     codeset.memory = std::move(program_image);
     process.LoadModule(std::move(codeset), process.GetEntryPoint());
@@ -210,8 +261,9 @@ static bool LoadNroImpl(Kernel::KProcess& process, const std::vector<u8>& data)
     return true;
 }
 
-bool AppLoader_NRO::LoadNro(Kernel::KProcess& process, const FileSys::VfsFile& nro_file) {
-    return LoadNroImpl(process, nro_file.ReadAllBytes());
+bool AppLoader_NRO::LoadNro(Core::System& system, Kernel::KProcess& process,
+                            const FileSys::VfsFile& nro_file) {
+    return LoadNroImpl(system, process, nro_file.ReadAllBytes());
 }
 
 AppLoader_NRO::LoadResult AppLoader_NRO::Load(Kernel::KProcess& process, Core::System& system) {
@@ -219,7 +271,7 @@ AppLoader_NRO::LoadResult AppLoader_NRO::Load(Kernel::KProcess& process, Core::S
         return {ResultStatus::ErrorAlreadyLoaded, {}};
     }
 
-    if (!LoadNro(process, *file)) {
+    if (!LoadNro(system, process, *file)) {
         return {ResultStatus::ErrorLoadingNRO, {}};
     }
 
diff --git a/src/core/loader/nro.h b/src/core/loader/nro.h
index 8de6eebc6..d2928cba0 100644
--- a/src/core/loader/nro.h
+++ b/src/core/loader/nro.h
@@ -54,7 +54,7 @@ public:
     bool IsRomFSUpdatable() const override;
 
 private:
-    bool LoadNro(Kernel::KProcess& process, const FileSys::VfsFile& nro_file);
+    bool LoadNro(Core::System& system, Kernel::KProcess& process, const FileSys::VfsFile& nro_file);
 
     std::vector<u8> icon_data;
     std::unique_ptr<FileSys::NACP> nacp;
diff --git a/src/core/loader/nso.cpp b/src/core/loader/nso.cpp
index 1350da8dc..34b10ef2e 100644
--- a/src/core/loader/nso.cpp
+++ b/src/core/loader/nso.cpp
@@ -20,6 +20,10 @@
 #include "core/loader/nso.h"
 #include "core/memory.h"
 
+#ifdef ARCHITECTURE_arm64
+#include "core/arm/nce/patch.h"
+#endif
+
 namespace Loader {
 namespace {
 struct MODHeader {
@@ -72,7 +76,8 @@ FileType AppLoader_NSO::IdentifyType(const FileSys::VirtualFile& in_file) {
 std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::System& system,
                                                const FileSys::VfsFile& nso_file, VAddr load_base,
                                                bool should_pass_arguments, bool load_into_process,
-                                               std::optional<FileSys::PatchManager> pm) {
+                                               std::optional<FileSys::PatchManager> pm,
+                                               Core::NCE::Patcher* patch) {
     if (nso_file.GetSize() < sizeof(NSOHeader)) {
         return std::nullopt;
     }
@@ -86,6 +91,16 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
         return std::nullopt;
     }
 
+    // Allocate some space at the beginning if we are patching in PreText mode.
+    const size_t module_start = [&]() -> size_t {
+#ifdef ARCHITECTURE_arm64
+        if (patch && patch->Mode() == Core::NCE::PatchMode::PreText) {
+            return patch->SectionSize();
+        }
+#endif
+        return 0;
+    }();
+
     // Build program image
     Kernel::CodeSet codeset;
     Kernel::PhysicalMemory program_image;
@@ -95,11 +110,12 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
         if (nso_header.IsSegmentCompressed(i)) {
             data = DecompressSegment(data, nso_header.segments[i]);
         }
-        program_image.resize(nso_header.segments[i].location + static_cast<u32>(data.size()));
-        std::memcpy(program_image.data() + nso_header.segments[i].location, data.data(),
-                    data.size());
-        codeset.segments[i].addr = nso_header.segments[i].location;
-        codeset.segments[i].offset = nso_header.segments[i].location;
+        program_image.resize(module_start + nso_header.segments[i].location +
+                             static_cast<u32>(data.size()));
+        std::memcpy(program_image.data() + module_start + nso_header.segments[i].location,
+                    data.data(), data.size());
+        codeset.segments[i].addr = module_start + nso_header.segments[i].location;
+        codeset.segments[i].offset = module_start + nso_header.segments[i].location;
         codeset.segments[i].size = nso_header.segments[i].size;
     }
 
@@ -118,7 +134,7 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
     }
 
     codeset.DataSegment().size += nso_header.segments[2].bss_size;
-    const u32 image_size{
+    u32 image_size{
         PageAlignSize(static_cast<u32>(program_image.size()) + nso_header.segments[2].bss_size)};
     program_image.resize(image_size);
 
@@ -139,6 +155,32 @@ std::optional<VAddr> AppLoader_NSO::LoadModule(Kernel::KProcess& process, Core::
         std::copy(pi_header.begin() + sizeof(NSOHeader), pi_header.end(), program_image.data());
     }
 
+#ifdef ARCHITECTURE_arm64
+    // If we are computing the process code layout and using nce backend, patch.
+    const auto& code = codeset.CodeSegment();
+    if (patch && patch->Mode() == Core::NCE::PatchMode::None) {
+        // Patch SVCs and MRS calls in the guest code
+        patch->PatchText(program_image, code);
+
+        // Add patch section size to the module size.
+        image_size += patch->SectionSize();
+    } 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->Mode() == Core::NCE::PatchMode::PreText ? 0 : image_size;
+        patch_segment.size = static_cast<u32>(patch->SectionSize());
+
+        // 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->Mode() == Core::NCE::PatchMode::PostData) {
+            image_size += patch_segment.size;
+        }
+    }
+#endif
+
     // If we aren't actually loading (i.e. just computing the process code layout), we are done
     if (!load_into_process) {
         return load_base + image_size;
diff --git a/src/core/loader/nso.h b/src/core/loader/nso.h
index 0b53b4ecd..29b86ed4c 100644
--- a/src/core/loader/nso.h
+++ b/src/core/loader/nso.h
@@ -15,6 +15,10 @@ namespace Core {
 class System;
 }
 
+namespace Core::NCE {
+class Patcher;
+}
+
 namespace Kernel {
 class KProcess;
 }
@@ -88,7 +92,8 @@ public:
     static std::optional<VAddr> LoadModule(Kernel::KProcess& process, Core::System& system,
                                            const FileSys::VfsFile& nso_file, VAddr load_base,
                                            bool should_pass_arguments, bool load_into_process,
-                                           std::optional<FileSys::PatchManager> pm = {});
+                                           std::optional<FileSys::PatchManager> pm = {},
+                                           Core::NCE::Patcher* patch = nullptr);
 
     LoadResult Load(Kernel::KProcess& process, Core::System& system) override;