From 3687b2e21477b53e59e6ce0401e3699b59c12373 Mon Sep 17 00:00:00 2001 From: MrMha <114929652+MrMha@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:49:37 +0400 Subject: [PATCH] perf(relinker): use a flat sorted call-site index --- core/relinker/CMakeLists.txt | 16 ++++++++ .../src/analysis/CallSiteResolver.cpp | 20 +++++++--- .../relinker/tests/CallSiteIndexTests.cpp | 37 +++++++++++++++++++ 3 files changed, 68 insertions(+), 5 deletions(-) create mode 100644 core/relinker/relinker/tests/CallSiteIndexTests.cpp diff --git a/core/relinker/CMakeLists.txt b/core/relinker/CMakeLists.txt index 095915729..4e8e1a4ad 100644 --- a/core/relinker/CMakeLists.txt +++ b/core/relinker/CMakeLists.txt @@ -161,6 +161,22 @@ if(BUILD_TESTING) add_test(NAME elf_reader_bounds COMMAND elf_reader_bounds_tests) set_tests_properties(elf_reader_bounds PROPERTIES TIMEOUT 20) + add_test_executable(call_site_index_tests + ${CMAKE_CURRENT_SOURCE_DIR}/relinker/tests/CallSiteIndexTests.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/CallSiteResolver.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/codegen/src/InstructionScanner.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/codegen/src/x86/X64InstructionDecoder.cpp + ) + target_include_directories(call_site_index_tests PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/domain/include + ${CMAKE_CURRENT_SOURCE_DIR}/relinker/include + ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/include + ${CMAKE_CURRENT_SOURCE_DIR}/codegen/include + ${CMAKE_CURRENT_SOURCE_DIR}/io/include + ) + add_test(NAME call_site_index COMMAND call_site_index_tests) + set_tests_properties(call_site_index PROPERTIES TIMEOUT 20) + add_test_executable(autorun_exit_tests ${CMAKE_CURRENT_SOURCE_DIR}/cli/tests/AutorunExitTests.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cli/src/Autorun.cpp diff --git a/core/relinker/relinker/src/analysis/CallSiteResolver.cpp b/core/relinker/relinker/src/analysis/CallSiteResolver.cpp index 6f08ab26f..0c698568e 100644 --- a/core/relinker/relinker/src/analysis/CallSiteResolver.cpp +++ b/core/relinker/relinker/src/analysis/CallSiteResolver.cpp @@ -4,7 +4,8 @@ #include #include #include -#include +#include +#include namespace Relinker { @@ -70,7 +71,7 @@ public: std::vector ResolveCallSites(const std::vector& textSection, FileByteOffset textSectionVAddr, VirtualAddress targetGotOrPltAddress, ByteCount targetGotOrPltSize) override; private: const std::vector* _cachedTextPtr = nullptr; - std::multimap _cachedTargetToInstr; + std::vector> _cachedTargetToInstr; FileByteOffset _cachedTextVAddr = 0; }; @@ -82,6 +83,10 @@ std::vector CallSiteResolver::ResolveCallSites(const std::vector auto scanner = Codegen::MakeInstructionScanner(); auto instructions = scanner->ScanCodeSection(textSection, 0, textSection.size()); _cachedTargetToInstr.clear(); + const auto referenceCount = std::count_if(instructions.begin(), instructions.end(), [&](const auto& match) { + return isRipRelativeFF(textSection, match) || isRipRelativeMov64(textSection, match); + }); + _cachedTargetToInstr.reserve(static_cast(referenceCount)); for (const auto& m : instructions) { std::size_t dispOff = 0; if (isRipRelativeFF(textSection, m)) @@ -94,15 +99,20 @@ std::vector CallSiteResolver::ResolveCallSites(const std::vector std::int32_t disp = readDisp32(textSection, dispOff); VirtualAddress instrEnd = textSectionVAddr + m.Offset + static_cast(m.Length); VirtualAddress target = static_cast(static_cast(instrEnd) + disp); - _cachedTargetToInstr.emplace(target, m.Offset); + _cachedTargetToInstr.emplace_back(target, m.Offset); } + std::sort(_cachedTargetToInstr.begin(), _cachedTargetToInstr.end()); _cachedTextPtr = &textSection; _cachedTextVAddr = textSectionVAddr; } std::vector sites; - auto lower = _cachedTargetToInstr.lower_bound(targetGotOrPltAddress); - auto upper = _cachedTargetToInstr.lower_bound(targetGotOrPltAddress + targetGotOrPltSize); + const auto findTarget = [&](VirtualAddress target) { + return std::lower_bound(_cachedTargetToInstr.begin(), _cachedTargetToInstr.end(), + std::pair{target, 0}); + }; + auto lower = findTarget(targetGotOrPltAddress); + auto upper = findTarget(targetGotOrPltAddress + targetGotOrPltSize); for (auto it = lower; it != upper; ++it) sites.push_back(textSectionVAddr + it->second); return sites; diff --git a/core/relinker/relinker/tests/CallSiteIndexTests.cpp b/core/relinker/relinker/tests/CallSiteIndexTests.cpp new file mode 100644 index 000000000..fceff1c26 --- /dev/null +++ b/core/relinker/relinker/tests/CallSiteIndexTests.cpp @@ -0,0 +1,37 @@ +#include +#include +#include +#include +#include + +static void Require(bool ok) { if (!ok) throw std::runtime_error("Incorrect call-site index"); } + +int main() { + try { + std::vector code; + const std::vector targets{0x2000, 0x1000, 0x2000, 0x3000, 0x1008}; + for (const auto target : targets) { + const auto offset = code.size(); + code.insert(code.end(), {0xFF, 0x15, 0, 0, 0, 0}); + const auto disp = static_cast(target - (0x100 + offset + 6)); + std::memcpy(code.data() + offset + 2, &disp, 4); + } + auto resolver = Relinker::MakeCallSiteResolver(); + Require(resolver->ResolveCallSites(code, 0x100, 0x2000, 8) == std::vector{0x100, 0x10C}); + Require(resolver->ResolveCallSites(code, 0x100, 0x1000, 8) == std::vector{0x106}); + Require(resolver->ResolveCallSites(code, 0x100, 0x1000, 9) == std::vector{0x106, 0x118}); + Require(resolver->ResolveCallSites(code, 0x100, 0x2000, 0).empty()); + Require(resolver->ResolveCallSites(code, 0x100, 0x4000, 8).empty()); + auto copy = code; + Require(resolver->ResolveCallSites(copy, 0x100, 0x2000, 8) == std::vector{0x100, 0x10C}); + Require(resolver->ResolveCallSites(copy, 0x200, 0x2100, 8) == std::vector{0x200, 0x20C}); + std::vector mov{0x48, 0x8B, 0x05, 0xF9, 0x0F, 0, 0}; + Require(resolver->ResolveCallSites(mov, 0x1000, 0x2000, 8) == std::vector{0x1000}); + Require(resolver->ResolveCallSites({}, 0, 0, 1).empty()); + bool rejected = false; + try { resolver->ResolveCallSites(code, 0, std::numeric_limits::max(), 1); } + catch (const std::exception&) { rejected = true; } + Require(rejected); + std::cout << "10 call-site index cases passed\n"; + } catch (const std::exception& error) { std::cerr << error.what() << '\n'; return 1; } +}