feat(libs): let relinked Windows games opt into mimalloc

.gitmodules

@@ -31,6 +31,9 @@
 [submodule "3rdparty/tracy"]
 	path = 3rdparty/tracy
 	url = https://github.com/wolfpld/tracy.git
+[submodule "3rdparty/mimalloc"]
+	path = 3rdparty/mimalloc
+	url = https://github.com/microsoft/mimalloc.git
 [submodule "3rdparty/libjpeg-turbo"]
 	path = 3rdparty/libjpeg-turbo
 	url = https://github.com/libjpeg-turbo/libjpeg-turbo.git

CMakeLists.txt

@@ -300,6 +300,22 @@ add_subdirectory(3rdparty/libusb-cmake EXCLUDE_FROM_ALL)
 set_target_properties(usb-1.0 PROPERTIES POSITION_INDEPENDENT_CODE ON)
 
 include(cmake/ZArchive.cmake)
+if(WIN32)
+    set(MI_BUILD_SHARED ON CACHE BOOL "" FORCE)
+    set(MI_BUILD_STATIC OFF CACHE BOOL "" FORCE)
+    set(MI_BUILD_OBJECT OFF CACHE BOOL "" FORCE)
+    set(MI_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+    set(MI_OVERRIDE ON CACHE BOOL "" FORCE)
+    set(MI_WIN_REDIRECT ON CACHE BOOL "" FORCE)
+    set(MI_FREE_IS_CHECKED SHARED CACHE STRING "" FORCE)
+    add_subdirectory(3rdparty/mimalloc EXCLUDE_FROM_ALL)
+    if(MINGW)
+        target_compile_definitions(mimalloc PRIVATE MI_MINGW_UCRT64=1)
+        target_link_options(mimalloc PRIVATE -static
+                LINKER:--wrap=malloc,--wrap=calloc,--wrap=realloc,--wrap=free,--wrap=strdup,--wrap=_strdup)
+        target_sources(mimalloc PRIVATE ${CMAKE_SOURCE_DIR}/core/libs/allocator/RuntimeHeap.c)
+    endif()
+endif()
 
 add_subdirectory(core/shader/recompiler)
 add_subdirectory(core/Decoder/Jpeg)

core/libs/CMakeLists.txt

@@ -198,6 +198,29 @@ if(APS5_ENABLE_TRACY)
     list(APPEND PATCHED_OUTPUTS ${TRACY_CLIENT_COPIED})
 endif()
 
+if(TARGET mimalloc)
+    add_executable(pe_import_first ${CMAKE_CURRENT_SOURCE_DIR}/allocator/PeImportFirstMain.cpp ${CMAKE_CURRENT_SOURCE_DIR}/allocator/PeImportFirst.cpp)
+    target_compile_features(pe_import_first PRIVATE cxx_std_20)
+    if(MINGW)
+        target_compile_options(pe_import_first PRIVATE -fasynchronous-unwind-tables)
+        target_link_options(pe_import_first PRIVATE -static -static-libgcc -static-libstdc++)
+    endif()
+    add_custom_command(
+            OUTPUT ${PATCHED_DIR}/mimalloc.dll
+            COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE:mimalloc> ${PATCHED_DIR}/mimalloc.dll
+            COMMAND $<TARGET_FILE:pe_import_first> ${PATCHED_DIR}/mimalloc.dll mimalloc-redirect.dll
+            DEPENDS mimalloc pe_import_first
+            COMMENT "Copying mimalloc.dll beside the libraries, its redirection module imported first"
+    )
+    add_custom_command(
+            OUTPUT ${PATCHED_DIR}/mimalloc-redirect.dll
+            COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE_DIR:mimalloc>/mimalloc-redirect.dll ${PATCHED_DIR}/mimalloc-redirect.dll
+            DEPENDS mimalloc
+            COMMENT "Copying mimalloc-redirect.dll beside the libraries"
+    )
+    list(APPEND PATCHED_OUTPUTS ${PATCHED_DIR}/mimalloc.dll ${PATCHED_DIR}/mimalloc-redirect.dll)
+endif()
+
 add_custom_target(libs DEPENDS ${PATCHED_OUTPUTS})
 
 if(BUILD_TESTING)
@@ -219,6 +242,15 @@ if(BUILD_TESTING)
     add_test(NAME nid_buffer_validation COMMAND nid_buffer_validation_tests)
     set_tests_properties(nid_buffer_validation PROPERTIES TIMEOUT 10)
 
+    add_test_executable(pe_import_first_tests allocator/tests/PeImportFirstTests.cpp allocator/PeImportFirst.cpp)
+    target_include_directories(pe_import_first_tests PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
+    target_compile_features(pe_import_first_tests PRIVATE cxx_std_20)
+    if(MINGW)
+        target_compile_options(pe_import_first_tests PRIVATE -fasynchronous-unwind-tables)
+    endif()
+    add_test(NAME pe_import_first COMMAND pe_import_first_tests)
+    set_tests_properties(pe_import_first PROPERTIES TIMEOUT 10)
+
     find_package(Python3 QUIET COMPONENTS Interpreter)
     if(Python3_Interpreter_FOUND)
         add_test(NAME nid_gnu_hash COMMAND ${Python3_EXECUTABLE}

core/libs/allocator/PeImportFirst.cpp

@@ -0,0 +1,93 @@
+#include "PeImportFirst.hpp"
+
+#include <algorithm>
+#include <array>
+#include <cctype>
+#include <cstdint>
+#include <cstring>
+#include <fstream>
+#include <iterator>
+#include <stdexcept>
+
+namespace Allocator {
+
+namespace {
+
+std::uint32_t read32(const std::vector<unsigned char>& bytes, std::size_t offset) {
+    if (offset + 4 > bytes.size()) throw std::runtime_error("truncated image");
+    std::uint32_t value = 0;
+    std::memcpy(&value, bytes.data() + offset, sizeof(value));
+    return value;
+}
+
+std::uint16_t read16(const std::vector<unsigned char>& bytes, std::size_t offset) {
+    if (offset + 2 > bytes.size()) throw std::runtime_error("truncated image");
+    std::uint16_t value = 0;
+    std::memcpy(&value, bytes.data() + offset, sizeof(value));
+    return value;
+}
+
+std::string lower(std::string text) {
+    std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
+    return text;
+}
+
+}
+
+void MoveImportFirst(std::vector<unsigned char>& bytes, const std::string& dll) {
+    const auto pe = read32(bytes, 0x3c);
+    if (read32(bytes, pe) != 0x00004550u) throw std::runtime_error("not a PE image");
+    const auto sections = read16(bytes, pe + 6);
+    const auto optionalSize = read16(bytes, pe + 20);
+    const auto optional = pe + 24;
+    if (read16(bytes, optional) != 0x20bu) throw std::runtime_error("not a PE32+ image");
+    const auto importRva = read32(bytes, optional + 112 + 8);
+    if (importRva == 0) throw std::runtime_error("the image has no imports");
+    const auto sectionTable = optional + optionalSize;
+    const auto fileOffset = [&](std::uint32_t rva) -> std::size_t {
+        for (std::uint16_t index = 0; index < sections; ++index) {
+            const auto header = sectionTable + index * 40u;
+            const auto virtualSize = read32(bytes, header + 8);
+            const auto virtualAddress = read32(bytes, header + 12);
+            const auto rawSize = read32(bytes, header + 16);
+            const auto rawOffset = read32(bytes, header + 20);
+            if (rva >= virtualAddress && rva < virtualAddress + std::max(virtualSize, rawSize)) {
+                if (rva - virtualAddress >= rawSize) throw std::runtime_error("RVA outside the section's file data");
+                return rawOffset + (rva - virtualAddress);
+            }
+        }
+        throw std::runtime_error("RVA outside every section");
+    };
+    const auto directory = fileOffset(importRva);
+    constexpr std::size_t DescriptorBytes = 20;
+    std::vector<std::array<unsigned char, DescriptorBytes>> descriptors;
+    std::size_t found = SIZE_MAX;
+    for (std::size_t offset = directory;; offset += DescriptorBytes) {
+        if (offset + DescriptorBytes > bytes.size()) throw std::runtime_error("truncated import directory");
+        std::array<unsigned char, DescriptorBytes> descriptor{};
+        std::memcpy(descriptor.data(), bytes.data() + offset, DescriptorBytes);
+        if (std::all_of(descriptor.begin(), descriptor.end(), [](unsigned char value) { return value == 0; })) break;
+        const auto nameOffset = fileOffset(read32(bytes, offset + 12));
+        std::string name;
+        for (auto at = nameOffset; at < bytes.size() && bytes[at] != 0; ++at) name.push_back(static_cast<char>(bytes[at]));
+        if (lower(name) == lower(dll)) found = descriptors.size();
+        descriptors.push_back(descriptor);
+    }
+    if (found == SIZE_MAX) throw std::runtime_error("the image does not import " + dll);
+    std::rotate(descriptors.begin(), descriptors.begin() + static_cast<std::ptrdiff_t>(found), descriptors.begin() + static_cast<std::ptrdiff_t>(found) + 1);
+    for (std::size_t index = 0; index < descriptors.size(); ++index) std::memcpy(bytes.data() + directory + index * DescriptorBytes, descriptors[index].data(), DescriptorBytes);
+}
+
+void MoveImportFirstInFile(const std::string& path, const std::string& dll) {
+    std::vector<unsigned char> bytes;
+    {
+        std::ifstream input(path, std::ios::binary);
+        if (!input) throw std::runtime_error("cannot open " + path);
+        bytes.assign(std::istreambuf_iterator<char>(input), std::istreambuf_iterator<char>());
+    }
+    MoveImportFirst(bytes, dll);
+    std::ofstream output(path, std::ios::binary | std::ios::trunc);
+    if (!output || !output.write(reinterpret_cast<const char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()))) throw std::runtime_error("cannot write " + path);
+}
+
+}

core/libs/allocator/PeImportFirst.hpp

@@ -0,0 +1,14 @@
+#ifndef CORE_LIBS_ALLOCATOR_PEIMPORTFIRST_HPP
+#define CORE_LIBS_ALLOCATOR_PEIMPORTFIRST_HPP
+
+#include <string>
+#include <vector>
+
+namespace Allocator {
+
+void MoveImportFirst(std::vector<unsigned char>& image, const std::string& dll);
+void MoveImportFirstInFile(const std::string& path, const std::string& dll);
+
+}
+
+#endif

core/libs/allocator/PeImportFirstMain.cpp

@@ -0,0 +1,18 @@
+#include "PeImportFirst.hpp"
+
+#include <cstdio>
+#include <exception>
+
+int main(int argc, char** argv) {
+    if (argc != 3) {
+        std::fprintf(stderr, "usage: pe_import_first <file> <dll>\n");
+        return 2;
+    }
+    try {
+        Allocator::MoveImportFirstInFile(argv[1], argv[2]);
+    } catch (const std::exception& error) {
+        std::fprintf(stderr, "pe_import_first: %s: %s\n", argv[1], error.what());
+        return 1;
+    }
+    return 0;
+}

core/libs/allocator/RuntimeHeap.c

@@ -0,0 +1,38 @@
+#include <windows.h>
+
+#include <stdint.h>
+#include <string.h>
+
+void* __wrap_malloc(size_t bytes) {
+    return HeapAlloc(GetProcessHeap(), 0, bytes != 0 ? bytes : 1);
+}
+
+void* __wrap_calloc(size_t count, size_t size) {
+    if (size != 0 && count > SIZE_MAX / size) return NULL;
+    const size_t bytes = count * size;
+    return HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, bytes != 0 ? bytes : 1);
+}
+
+void __wrap_free(void* block) {
+    if (block != NULL) HeapFree(GetProcessHeap(), 0, block);
+}
+
+void* __wrap_realloc(void* block, size_t bytes) {
+    if (block == NULL) return __wrap_malloc(bytes);
+    if (bytes == 0) {
+        __wrap_free(block);
+        return NULL;
+    }
+    return HeapReAlloc(GetProcessHeap(), 0, block, bytes);
+}
+
+char* __wrap__strdup(const char* text) {
+    const size_t bytes = strlen(text) + 1;
+    char* copy = __wrap_malloc(bytes);
+    if (copy != NULL) memcpy(copy, text, bytes);
+    return copy;
+}
+
+char* __wrap_strdup(const char* text) {
+    return __wrap__strdup(text);
+}

core/libs/allocator/tests/PeImportFirstTests.cpp

@@ -0,0 +1,156 @@
+#include "allocator/PeImportFirst.hpp"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <iterator>
+#include <random>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace {
+
+constexpr std::size_t PeOffset = 0x40;
+constexpr std::size_t OptionalOffset = PeOffset + 24;
+constexpr std::size_t OptionalBytes = 240;
+constexpr std::size_t SectionTable = OptionalOffset + OptionalBytes;
+constexpr std::uint32_t SectionRva = 0x1000;
+constexpr std::size_t SectionOffset = 0x200;
+constexpr std::size_t SectionBytes = 0x200;
+constexpr std::size_t DescriptorBytes = 20;
+
+void require(bool condition, const std::string& message) {
+    if (!condition) throw std::runtime_error(message);
+}
+
+void put16(std::vector<unsigned char>& bytes, std::size_t offset, std::uint16_t value) {
+    std::memcpy(bytes.data() + offset, &value, sizeof(value));
+}
+
+void put32(std::vector<unsigned char>& bytes, std::size_t offset, std::uint32_t value) {
+    std::memcpy(bytes.data() + offset, &value, sizeof(value));
+}
+
+std::uint32_t get32(const std::vector<unsigned char>& bytes, std::size_t offset) {
+    std::uint32_t value = 0;
+    std::memcpy(&value, bytes.data() + offset, sizeof(value));
+    return value;
+}
+
+std::vector<unsigned char> image(const std::vector<std::string>& dlls, std::uint16_t magic = 0x20b) {
+    std::vector<unsigned char> bytes(SectionOffset + SectionBytes, 0);
+    bytes[0] = 'M';
+    bytes[1] = 'Z';
+    put32(bytes, 0x3c, PeOffset);
+    put32(bytes, PeOffset, 0x00004550u);
+    put16(bytes, PeOffset + 4, 0x8664);
+    put16(bytes, PeOffset + 6, 1);
+    put16(bytes, PeOffset + 20, OptionalBytes);
+    put16(bytes, OptionalOffset, magic);
+    put32(bytes, OptionalOffset + 108, 16);
+    put32(bytes, OptionalOffset + 112 + 8, SectionRva);
+    put32(bytes, OptionalOffset + 112 + 12, static_cast<std::uint32_t>((dlls.size() + 1) * DescriptorBytes));
+    std::memcpy(bytes.data() + SectionTable, ".idata", 6);
+    put32(bytes, SectionTable + 8, SectionBytes);
+    put32(bytes, SectionTable + 12, SectionRva);
+    put32(bytes, SectionTable + 16, SectionBytes);
+    put32(bytes, SectionTable + 20, SectionOffset);
+    auto name = (dlls.size() + 1) * DescriptorBytes;
+    for (std::size_t index = 0; index < dlls.size(); ++index) {
+        const auto descriptor = SectionOffset + index * DescriptorBytes;
+        put32(bytes, descriptor, static_cast<std::uint32_t>(SectionRva + 0x100 + index * 0x10));
+        put32(bytes, descriptor + 4, static_cast<std::uint32_t>(index + 1));
+        put32(bytes, descriptor + 12, static_cast<std::uint32_t>(SectionRva + name));
+        put32(bytes, descriptor + 16, static_cast<std::uint32_t>(SectionRva + 0x180 + index * 0x10));
+        std::memcpy(bytes.data() + SectionOffset + name, dlls[index].c_str(), dlls[index].size() + 1);
+        name += dlls[index].size() + 1;
+    }
+    return bytes;
+}
+
+std::string importName(const std::vector<unsigned char>& bytes, std::size_t index) {
+    const auto rva = get32(bytes, SectionOffset + index * DescriptorBytes + 12);
+    return reinterpret_cast<const char*>(bytes.data() + SectionOffset + (rva - SectionRva));
+}
+
+std::vector<unsigned char> descriptor(const std::vector<unsigned char>& bytes, std::size_t index) {
+    const auto first = bytes.begin() + static_cast<std::ptrdiff_t>(SectionOffset + index * DescriptorBytes);
+    return {first, first + static_cast<std::ptrdiff_t>(DescriptorBytes)};
+}
+
+void writeFile(const std::filesystem::path& path, const std::vector<unsigned char>& bytes) {
+    std::ofstream output(path, std::ios::binary | std::ios::trunc);
+    output.write(reinterpret_cast<const char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()));
+    require(static_cast<bool>(output), "cannot write " + path.string());
+}
+
+std::vector<unsigned char> readFile(const std::filesystem::path& path) {
+    std::ifstream input(path, std::ios::binary);
+    require(static_cast<bool>(input), "cannot read " + path.string());
+    return {std::istreambuf_iterator<char>(input), std::istreambuf_iterator<char>()};
+}
+
+void movesTheNamedImportFirst() {
+    const auto original = image({"kernel32.dll", "ucrtbase.dll", "mimalloc-redirect.dll", "advapi32.dll"});
+    auto moved = original;
+    Allocator::MoveImportFirst(moved, "MIMALLOC-REDIRECT.dll");
+    require(importName(moved, 0) == "mimalloc-redirect.dll" && importName(moved, 1) == "kernel32.dll" && importName(moved, 2) == "ucrtbase.dll" && importName(moved, 3) == "advapi32.dll",
+        "the named import did not move first, or the others changed their order");
+    require(descriptor(moved, 0) == descriptor(original, 2) && descriptor(moved, 1) == descriptor(original, 0) && descriptor(moved, 2) == descriptor(original, 1) && descriptor(moved, 3) == descriptor(original, 3),
+        "a descriptor did not move whole");
+    const auto directoryEnd = SectionOffset + 4 * DescriptorBytes;
+    require(std::equal(moved.begin(), moved.begin() + SectionOffset, original.begin()) && std::equal(moved.begin() + directoryEnd, moved.end(), original.begin() + directoryEnd),
+        "bytes outside the import descriptors changed");
+    auto unchanged = original;
+    Allocator::MoveImportFirst(unchanged, "kernel32.dll");
+    require(unchanged == original, "moving the first import changed the image");
+}
+
+void movesTheImportInAFile(const std::filesystem::path& directory) {
+    const auto path = directory / "moved.dll";
+    writeFile(path, image({"kernel32.dll", "mimalloc-redirect.dll"}));
+    Allocator::MoveImportFirstInFile(path.string(), "mimalloc-redirect.dll");
+    const auto moved = readFile(path);
+    require(importName(moved, 0) == "mimalloc-redirect.dll" && importName(moved, 1) == "kernel32.dll", "the file's import did not move first");
+}
+
+void failsWithoutTouchingTheFile(const std::filesystem::path& directory, const std::string& what, const std::vector<unsigned char>& bytes, const std::string& expected) {
+    const auto path = directory / (what + ".dll");
+    writeFile(path, bytes);
+    try {
+        Allocator::MoveImportFirstInFile(path.string(), "mimalloc-redirect.dll");
+    } catch (const std::exception& error) {
+        require(std::string(error.what()).find(expected) != std::string::npos, what + ": failed with another error: " + error.what());
+        require(readFile(path) == bytes, what + ": the file changed although the move failed");
+        return;
+    }
+    throw std::runtime_error(what + ": the move did not fail");
+}
+
+}
+
+int main() {
+    const auto directory = std::filesystem::temp_directory_path() / ("aps5-pe-import-first-" + std::to_string(std::random_device{}()));
+    try {
+        std::filesystem::create_directories(directory);
+        movesTheNamedImportFirst();
+        movesTheImportInAFile(directory);
+        failsWithoutTouchingTheFile(directory, "missing-import", image({"kernel32.dll", "ucrtbase.dll"}), "does not import mimalloc-redirect.dll");
+        auto truncated = image({"kernel32.dll", "mimalloc-redirect.dll", "ucrtbase.dll"});
+        truncated.resize(SectionOffset + DescriptorBytes + DescriptorBytes / 2);
+        failsWithoutTouchingTheFile(directory, "truncated-directory", truncated, "truncated import directory");
+        failsWithoutTouchingTheFile(directory, "pe32", image({"kernel32.dll", "mimalloc-redirect.dll"}, 0x10b), "not a PE32+ image");
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        std::error_code ignored;
+        std::filesystem::remove_all(directory, ignored);
+        return 1;
+    }
+    std::filesystem::remove_all(directory);
+    std::cout << "pe import first tests passed\n";
+    return 0;
+}

core/relinker/CMakeLists.txt

@@ -659,14 +659,8 @@ if(BUILD_TESTING)
     )
     add_test(NAME windows_icon_resource COMMAND windows_icon_resource_tests)
 
-    add_test_executable(entry_stub_tests
-        ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/tests/EntryStubTests.cpp
-        ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/general/EntryStubBuilder.cpp
-    )
-    target_include_directories(entry_stub_tests PRIVATE
-        ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/include
-        ${CMAKE_CURRENT_SOURCE_DIR}/domain/include
-    )
+    add_test_executable(entry_stub_tests ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/tests/EntryStubTests.cpp)
+    target_link_libraries(entry_stub_tests PRIVATE relinker_core)
     add_test(NAME entry_stub COMMAND entry_stub_tests)
 
     add_test_executable(control_flow_graph_tests

core/relinker/elfpatcher/src/windows/WindowsEntryStubBuilder.cpp

@@ -98,6 +98,7 @@ WindowsEntryStub WindowsEntryStubBuilder::Build(const std::uint32_t dataRva, con
     };
 
     const auto parentVariable = addWideString(u"ANYPS5_LOADEXEC_PARENT");
+    const auto allocatorVariable = addWideString(u"APS5_MIMALLOC");
     const auto parentText = reserve(34);
     const auto skipFirstArgument = reserve(4);
     const auto programPath = reserve(PathCapacity * 2);
@@ -119,6 +120,17 @@ WindowsEntryStub WindowsEntryStubBuilder::Build(const std::uint32_t dataRva, con
     const auto unwindRva = CheckedRva(dataRva + data.size());
     data.insert(data.end(), {1, 10, 6, 0, 10, 0xb2, 6, 0xc0, 4, 0x70, 3, 0x60, 2, 0x50, 1, 0x30});
 
+    std::vector<std::uint32_t> allocatorPaths;
+    std::vector<std::size_t> allocatorUnits;
+    for (const std::string library : {"mimalloc.dll"}) {
+        const auto name = path + library;
+        const auto wideName = toUtf16(name);
+        if (wideName.size() + 1 >= PathCapacity)
+            throw Domain::RelinkerException("Windows library path exceeds the startup buffer: " + name);
+        allocatorPaths.push_back(addWideString(wideName));
+        allocatorUnits.push_back(wideName.size());
+    }
+
     std::vector<std::uint32_t> libraryPaths;
     std::vector<std::size_t> libraryUnits;
     for (std::size_t index = 0; index < libraries.size(); ++index) {
@@ -479,6 +491,35 @@ WindowsEntryStub WindowsEntryStubBuilder::Build(const std::uint32_t dataRva, con
         writeUtf8(programText);
     }
 
+    code.Rip({0x48, 0x8d, 0x0d}, allocatorVariable);
+    code.Emit({0x31, 0xd2, 0x45, 0x31, 0xc0});
+    call("GetEnvironmentVariableW");
+    code.Emit({0x85, 0xc0});
+    const auto noAllocator = code.Branch({0x0f, 0x84});
+    for (std::size_t index = 0; index < allocatorPaths.size(); ++index) {
+        std::optional<std::size_t> tooLong;
+        if (absolutePath) {
+            code.Rip({0x48, 0x8d, 0x0d}, allocatorPaths[index]);
+        } else {
+            const auto units = CheckedRva(allocatorUnits[index] + 1);
+            code.Emit({0x4c, 0x89, 0xe0, 0x48, 0x29, 0xd8, 0x48, 0x05});
+            code.U32(units * 2);
+            code.Emit({0x48, 0x3d});
+            code.U32(PathCapacity * 2);
+            tooLong = code.Branch({0x0f, 0x87});
+            code.Emit({0x4c, 0x89, 0xe7});
+            code.Rip({0x48, 0x8d, 0x35}, allocatorPaths[index]);
+            code.Emit({0xb9});
+            code.U32(units);
+            code.Emit({0x66, 0xf3, 0xa5});
+            code.Rip({0x48, 0x8d, 0x0d}, modulePath);
+        }
+        code.Emit({0x31, 0xd2, 0x41, 0xb8, 0, 0x11, 0, 0});
+        call("LoadLibraryExW");
+        if (tooLong) code.PatchBranch(*tooLong, code.GetRva());
+    }
+    code.PatchBranch(noAllocator, code.GetRva());
+
     for (std::size_t index = 0; index < libraries.size(); ++index) {
         if (index >= guestModules.size() && !imported[index] && !unscopedImports) continue;
         if (absolutePath && index >= guestModules.size()) {

core/relinker/elfpatcher/tests/EntryStubTests.cpp

@@ -1,9 +1,12 @@
 #include <elfpatcher/general/EntryStubBuilder.hpp>
+#include <elfpatcher/windows/WindowsEntryStubBuilder.hpp>
+#include <domain/Types.hpp>
 #include <cstdint>
 #include <cstring>
 #include <iostream>
 #include <stdexcept>
 #include <string>
+#include <string_view>
 #include <vector>
 
 namespace {
@@ -36,11 +39,93 @@ void stubTerminatesTheFrameChain() {
     require(bytesAt(stub, offset, {0x0f, 0x0b}) && offset == stub.size(), "the stub does not end with ud2 after the call");
 }
 
+std::uint32_t wideStringRva(const Elfpatcher::Windows::PeSection& section, const std::u16string& text) {
+    std::vector<std::uint8_t> needle;
+    for (const auto unit : text) {
+        needle.push_back(static_cast<std::uint8_t>(unit));
+        needle.push_back(static_cast<std::uint8_t>(unit >> 8));
+    }
+    needle.insert(needle.end(), {0, 0});
+    for (std::size_t offset = 0; offset + needle.size() <= section.Data.size(); offset += 2) {
+        if (std::memcmp(section.Data.data() + offset, needle.data(), needle.size()) == 0) return section.Rva + static_cast<std::uint32_t>(offset);
+    }
+    throw std::runtime_error("the stub data does not hold the expected path");
+}
+
+std::size_t ripReference(const Elfpatcher::Windows::PeSection& code, const std::vector<std::uint8_t>& opcode, const std::uint32_t target, const std::size_t from = 0) {
+    for (std::size_t offset = from; offset + opcode.size() + 4 <= code.Data.size(); ++offset) {
+        if (std::memcmp(code.Data.data() + offset, opcode.data(), opcode.size()) != 0) continue;
+        std::int32_t displacement = 0;
+        std::memcpy(&displacement, code.Data.data() + offset + opcode.size(), sizeof(displacement));
+        if (code.Rva + offset + opcode.size() + 4 + static_cast<std::int64_t>(displacement) == target) return offset;
+    }
+    return std::string::npos;
+}
+
+Elfpatcher::Windows::WindowsEntryStub windowsStub(const Elfpatcher::Windows::WindowsImports& nativeImports, const std::string& runPath, const std::vector<std::string>& libraries) {
+    std::vector<Elfpatcher::Windows::PeImport> imports;
+    for (std::size_t index = 0; index < libraries.size(); ++index) {
+        imports.push_back({.Name = "import" + std::to_string(index), .TargetRva = static_cast<std::uint32_t>(0x2000 + index * 8), .Addend = 0, .Library = libraries[index]});
+    }
+    return Elfpatcher::Windows::WindowsEntryStubBuilder().Build(0x40000, 0x1000, nativeImports, libraries, imports, runPath, false, false);
+}
+
+void windowsStubLoadsTheAllocatorFirst(const std::string& runPath, const std::u16string& directory, const std::vector<std::uint8_t>& pathOpcode) {
+    const auto nativeImports = Elfpatcher::Windows::WindowsImportBuilder().Build(0x1000);
+    const auto stub = windowsStub(nativeImports, runPath, {"libc.prx", "libkernel.prx"});
+    const auto allocator = ripReference(stub.Code, pathOpcode, wideStringRva(stub.Data, directory + u"mimalloc.dll"));
+    const auto library = ripReference(stub.Code, pathOpcode, wideStringRva(stub.Data, directory + u"libc.prx"));
+    require(allocator != std::string::npos, runPath + ": the stub does not use mimalloc.dll's path");
+    require(library != std::string::npos, runPath + ": the stub does not use the first library's path");
+    const auto load = ripReference(stub.Code, {0xff, 0x15}, nativeImports.Functions.at("LoadLibraryExW"), allocator);
+    require(allocator < library && load < library, runPath + ": the stub loads a library before mimalloc.dll");
+    std::size_t flags = load - 8;
+    require(bytesAt(stub.Code.Data, flags, {0x31, 0xd2, 0x41, 0xb8, 0x00, 0x11, 0x00, 0x00}), runPath + ": mimalloc.dll is not loaded from the stub's search directories");
+}
+
+void windowsStubLoadsTheAllocatorOnRequest(const std::string& runPath, const std::u16string& directory, const std::vector<std::uint8_t>& pathOpcode) {
+    const auto nativeImports = Elfpatcher::Windows::WindowsImportBuilder().Build(0x1000);
+    const auto stub = windowsStub(nativeImports, runPath, {"libc.prx"});
+    const auto variable = ripReference(stub.Code, {0x48, 0x8d, 0x0d}, wideStringRva(stub.Data, u"APS5_MIMALLOC"));
+    const auto allocator = ripReference(stub.Code, pathOpcode, wideStringRva(stub.Data, directory + u"mimalloc.dll"));
+    require(variable != std::string::npos && allocator != std::string::npos && variable < allocator, runPath + ": the stub does not read APS5_MIMALLOC before loading mimalloc.dll");
+    std::size_t cursor = variable + 7;
+    require(bytesAt(stub.Code.Data, cursor, {0x31, 0xd2, 0x45, 0x31, 0xc0}) && ripReference(stub.Code, {0xff, 0x15}, nativeImports.Functions.at("GetEnvironmentVariableW"), cursor) == cursor, runPath + ": APS5_MIMALLOC is not read with GetEnvironmentVariableW");
+    cursor += 6;
+    require(bytesAt(stub.Code.Data, cursor, {0x85, 0xc0, 0x0f, 0x84}), runPath + ": an unset APS5_MIMALLOC does not branch");
+    std::int32_t displacement = 0;
+    std::memcpy(&displacement, stub.Code.Data.data() + cursor, sizeof(displacement));
+    const auto target = static_cast<std::int64_t>(cursor) + 4 + displacement;
+    const auto load = ripReference(stub.Code, {0xff, 0x15}, nativeImports.Functions.at("LoadLibraryExW"), allocator);
+    require(load != std::string::npos && target == static_cast<std::int64_t>(load) + 6, runPath + ": an unset APS5_MIMALLOC does not skip exactly the load of mimalloc.dll");
+}
+
+void windowsStubRejectsAnOverlongAllocatorPath() {
+    const auto nativeImports = Elfpatcher::Windows::WindowsImportBuilder().Build(0x1000);
+    const auto runPath = [](const std::size_t length) { return "C:\\" + std::string(length - 4, 'a') + "\\"; };
+    constexpr std::size_t PathUnits = 32768;
+    const auto longest = PathUnits - 1 - std::string_view("mimalloc.dll").size() - 1;
+    static_cast<void>(windowsStub(nativeImports, runPath(longest), {"libc.prx"}));
+    try {
+        static_cast<void>(windowsStub(nativeImports, runPath(longest + 1), {"libc.prx"}));
+    } catch (const Domain::RelinkerException& error) {
+        const std::string message = error.what();
+        require(message.starts_with("Windows library path exceeds the startup buffer: ") && message.ends_with("\\mimalloc.dll"), "an overlong allocator path fails with another error: " + message);
+        return;
+    }
+    throw std::runtime_error("an allocator path past the startup buffer was accepted");
+}
+
 }
 
 int main() {
     try {
         stubTerminatesTheFrameChain();
+        windowsStubLoadsTheAllocatorFirst("C:/Games/Title/libs", u"C:\\Games\\Title\\libs\\", {0x48, 0x8d, 0x0d});
+        windowsStubLoadsTheAllocatorFirst("$ORIGIN/libs", u"libs\\", {0x48, 0x8d, 0x35});
+        windowsStubLoadsTheAllocatorOnRequest("C:/Games/Title/libs", u"C:\\Games\\Title\\libs\\", {0x48, 0x8d, 0x0d});
+        windowsStubLoadsTheAllocatorOnRequest("$ORIGIN/libs", u"libs\\", {0x48, 0x8d, 0x35});
+        windowsStubRejectsAnOverlongAllocatorPath();
     } catch (const std::exception& error) {
         std::cerr << error.what() << '\n';
         return 1;

docs/dev/ARCHITECTURE.md

@@ -53,6 +53,19 @@ While initialization runs, the count is `UINT32_MAX` and the entry word holds th
 
 [GuestDeferredModule.cpp](../../core/libs/tests/GuestDeferredModule.cpp) checks repeated loads at the same base, multiple guest references, an external native reference, concurrent initialization and recursive loading. The Windows fixture uses the same table layout as the relinker and reports its entry and constructor calls separately.
 
+## Windows allocator
+
+On Windows, a relinked game run with `APS5_MIMALLOC` set (to any value) allocates through [mimalloc](../../3rdparty/mimalloc) instead of the C runtime's heap: its libraries and the C and C++ runtime they use. Without the variable the game keeps the C runtime's heap.
+
+- The build copies `mimalloc.dll` and `mimalloc-redirect.dll` beside the libraries ([`core/libs/CMakeLists.txt`](../../core/libs/CMakeLists.txt)). [`pe_import_first`](../../core/libs/allocator/PeImportFirst.cpp) moves `mimalloc-redirect.dll` to the front of the copied `mimalloc.dll`'s import directory. The redirection module has to load before mimalloc, and MinGW's linker does not keep that order.
+- When `APS5_MIMALLOC` is set, the relinker's Windows entry stub ([`WindowsEntryStubBuilder.cpp`](../../core/relinker/elfpatcher/src/windows/WindowsEntryStubBuilder.cpp)) loads `mimalloc.dll` from the run path with `LoadLibraryExW` before the first library. `mimalloc-redirect.dll` then sends the C runtime's allocation functions in `ucrtbase.dll` (`malloc`, `calloc`, `realloc`, `free` and their relatives) to mimalloc for the whole process. The libraries loaded after it, and the MinGW runtime DLLs they use (`libstdc++`, `libgcc_s`, `libwinpthread`), allocate and free through mimalloc.
+  - Without `mimalloc.dll` the load fails quietly and the process keeps the C runtime's heap. With a relative run path, a module directory too long for the stub's path buffer skips the load the same way.
+  - A run path whose `mimalloc.dll` path cannot fit that buffer fails the relink, as a library path does.
+- Guest code, the executable and its bundled modules, allocates through libc's guest heap ([`GuestHeap.cpp`](../../core/libs/prx/libc/src/GuestHeap.cpp)). The guest heap takes its blocks from the guest arena, or from `malloc` when there is no arena, and returns them the same way, so a block never moves between the guest heap and the host allocator.
+- `mimalloc.dll` links the MinGW runtime statically, and three parts of it allocate through the C runtime: GCC's emulated thread-locals (`emutls`, which hold mimalloc's own thread-locals), MinGW's thread-exit destructor list and winpthreads.
+  - After the redirection these calls would enter mimalloc again, before the calling thread's mimalloc state exists.
+  - Linker wraps send them to [`RuntimeHeap.c`](../../core/libs/allocator/RuntimeHeap.c), which allocates from the process heap (`HeapAlloc` on `GetProcessHeap()`). The same code frees these blocks through the same wraps.
+
 ## Thread cancellation
 
 [`Thread.cpp`](../../core/libs/prx/libkernel/Pthread/src/Thread.cpp) records cancellation requests on the target guest thread. New threads start with cancellation enabled and deferred. A disabled request stays pending; enabling deferred cancellation returns before the next cancellation point acts on it. Repeated requests do not add cleanup invocations.

docs/dev/ENVIRONMENT.md

@@ -170,6 +170,7 @@ The code reads environment names through `getenv`, `IsEnvironmentSet`, `kibSetti
 | `APS5_LOG_STATS` | presence | `core/libs/prx/libSceAgcDriver/Execution/include/AsyncLog.hpp` |  | prints the AGC log's line, byte and coalescing statistics every window; 0 turns them off |
 | `APS5_LOG_SYNC` | presence | `core/libs/prx/libSceAgcDriver/Execution/include/AsyncLog.hpp` |  | writes AGC log lines on the calling thread instead of through the asynchronous sink thread |
 | `APS5_LOOP_GUARD` | presence | `core/libs/prx/libSceAgcDriver/Graphics/src/ShaderResources.cpp` |  | enables the fault buffer that guards shader loops running past their bound |
+| `APS5_MIMALLOC` | presence | `core/relinker/elfpatcher/src/windows/WindowsEntryStubBuilder.cpp` |  | loads `mimalloc.dll` in the Windows entry stub, so the process allocates through mimalloc instead of the C runtime's heap |
 | `APS5_MIRROR_REFRESH_THREADS` | integer | `core/libs/prx/libSceAgcDriver/Graphics/src/GuestBufferMemory.cpp` | 3 | number of helper threads that refresh heap mirrors, clamped to 0 through 16 |
 | `APS5_NO_ADDRESS_READ_PASSES` | presence | `core/libs/prx/libSceAgcDriver/Graphics/src/Draw.cpp` |  | stops a draw that only reads through its buffer descriptors from counting as read-only, so it breaks the pass instead of continuing the producer's |
 | `APS5_NO_ADDRESS_SPACE_CACHE` | presence | `core/libs/prx/libSceAgcDriver/Graphics/src/GuestBufferMemory.cpp` |  | disables the address-space cache so every build rebuilds its space |

tools/package_release.py

@@ -49,6 +49,7 @@ def package(platform, build, output, version):
     if platform == "windows":
         runtime = windows_runtime(build)
         files.extend(runtime / name for name in ("libgcc_s_seh-1.dll", "libstdc++-6.dll", "libwinpthread-1.dll"))
+        files.extend(build / "core/libs/libs" / name for name in ("mimalloc.dll", "mimalloc-redirect.dll"))
     files.extend(tracy_runtime(platform, build))
     binary = build / "core/relinker" / executable
     for file in [*files, binary]:

tools/tests/test_package_release.py

@@ -53,16 +53,19 @@ class ReleasePackageTests(unittest.TestCase):
         names = ("libgcc_s_seh-1.dll", "libstdc++-6.dll", "libwinpthread-1.dll")
         for name in names:
             (runtime / name).write_bytes(name.encode())
+        allocator = ("mimalloc.dll", "mimalloc-redirect.dll")
+        for name in allocator:
+            (self.library.parent / name).write_bytes(name.encode())
         compiler = runtime / "g++.exe"
         (self.build / "CMakeCache.txt").write_text("CMAKE_CXX_COMPILER:STRING=" + compiler.as_posix() + "\n", encoding="utf-8")
         (self.build / "core/relinker/relinker.exe").write_bytes(b"Windows executable")
         package_release.package("windows", self.build, self.output, "v1")
         with zipfile.ZipFile(self.output / "prx-windows-v1.zip") as archive:
-            self.assertEqual(set(archive.namelist()), {"libs/libc.prx", *("libs/" + name for name in names)})
-            for name in names:
+            self.assertEqual(set(archive.namelist()), {"libs/libc.prx", *("libs/" + name for name in names + allocator)})
+            for name in names + allocator:
                 self.assertEqual(archive.read("libs/" + name), name.encode())
         with tarfile.open(self.output / "prx-windows-v1.tar.gz") as archive:
-            for name in names:
+            for name in names + allocator:
                 self.assertEqual(archive.extractfile("libs/" + name).read(), name.encode())
         self.assertEqual((self.output / "relinker-v1.exe").read_bytes(), b"Windows executable")
 
@@ -91,6 +94,9 @@ class ReleasePackageTests(unittest.TestCase):
                     for runtime_name in ("libgcc_s_seh-1.dll", "libstdc++-6.dll", "libwinpthread-1.dll"):
                         (runtime / runtime_name).write_bytes(runtime_name.encode())
                         expected.add("libs/" + runtime_name)
+                    for allocator_name in ("mimalloc.dll", "mimalloc-redirect.dll"):
+                        (directory / allocator_name).write_bytes(allocator_name.encode())
+                        expected.add("libs/" + allocator_name)
                     cache += "CMAKE_CXX_COMPILER:FILEPATH=" + (runtime / "g++.exe").as_posix() + "\n"
                     (self.build / "core/relinker/relinker.exe").write_bytes(b"Windows executable")
                 (self.build / "CMakeCache.txt").write_text(cache, encoding="utf-8")