// SPDX-License-Identifier: MIT
#include "FEXCore/Utils/LogManager.h"
#include "FEXCore/Utils/MathUtils.h"
#include "FEXCore/Utils/TypeDefines.h"
#include "FEXCore/fextl/memory.h"
#include <FEXCore/Utils/Profiler.h>
#include <FEXCore/Utils/SpinWaitLock.h>

#include <Interface/Context/Context.h>
#include <Interface/Core/ArchHelpers/Arm64Emitter.h>
#include <Interface/Core/Dispatcher/Dispatcher.h>
#include <Interface/Core/JIT/DebugData.h>
#include <Interface/Core/JIT/Relocations.h>
#include <Interface/Core/LookupCache.h>
#include <Interface/Core/OpcodeDispatcher.h>
#include <Interface/IR/PassManager.h>

#include <FEXCore/Core/Thunks.h>
#include <FEXCore/HLE/SourcecodeResolver.h>
#include <FEXCore/HLE/SyscallHandler.h>

#include <FEXHeaderUtils/Filesystem.h>

#include <algorithm>
#include <git_version.h>

#include <span>
#include <xxhash.h>

#include <FEXCore/Utils/AllocatorHooks.h>

#include <fstream>

namespace FEXCore {

#if __clang_major__ < 16
ExecutableFileInfo::ExecutableFileInfo(fextl::unique_ptr<HLE::SourcecodeMap> Map, uint64_t FileId, fextl::string Filename)
  : SourcecodeMap(std::move(Map))
  , FileId(FileId)
  , Filename(Filename) {}
#endif
ExecutableFileInfo::~ExecutableFileInfo() = default;

MappedCodeCacheFile::~MappedCodeCacheFile() {
  if (CacheManager) {
    CacheManager->UnregisterMappedCodeBuffer(*this);
  }

#ifndef _WIN32
  if (!CodeBuffer.empty()) {
    FEXCore::Allocator::munmap(CodeBuffer.data(), CodeBuffer.size_bytes());
  }
#endif
}

void AbstractCodeCache::RegisterMappedCodeBuffer(MappedCodeCacheFile& Code) {
  MappedCodeBuffers.push_back(Code.CodeBuffer);
  // Unregister on destruction of Code
  Code.CacheManager = this;
}

void AbstractCodeCache::UnregisterMappedCodeBuffer(MappedCodeCacheFile& Code) {
  std::erase_if(MappedCodeBuffers, [&](const auto& Elem) { return Elem.data() == Code.CodeBuffer.data(); });
}

bool AbstractCodeCache::IsAddressInMappedCodeBuffer(uintptr_t Address) const {
  for (const auto& Range : MappedCodeBuffers) {
    auto Start = reinterpret_cast<uintptr_t>(Range.data());
    if (Address >= Start && Address < Start + Range.size_bytes()) {
      return true;
    }
  }
  return false;
}

fextl::string CodeMap::GetBaseFilename(const ExecutableFileInfo& MainExecutable, bool AddNombSuffix) {
  auto FileId = MainExecutable.FileId;

  std::string_view base_filename = FHU::Filesystem::GetFilename(std::string_view {MainExecutable.Filename});
  if (FileId != 0xffff'ffff'ffff'ffff) {
    return fextl::fmt::format("{}-{:016x}{}", base_filename, MainExecutable.FileId, AddNombSuffix ? "-nomb" : "");
  }

  return "";
}

fextl::map<CodeMapFileId, CodeMap::ParsedContents> CodeMap::ParseCodeMap(std::ifstream& File) {
  fextl::map<CodeMapFileId, CodeMap::ParsedContents> Ret;
  while (true) {
    Entry Entry;
    File.read(reinterpret_cast<char*>(&Entry), sizeof(Entry));
    if (!File) {
      break;
    }

    if (Entry.FileId == LoadExternalLibrary.FileId && Entry.BlockOffset == LoadExternalLibrary.BlockOffset) {
      ExternalLibraryInfo Info;
      File.read(reinterpret_cast<char*>(&Info), sizeof(Info));

      fextl::string Filename;
      std::getline(File, Filename, '\0');

      // Align to 4-byte boundary
      char Null[4];
      File.read(Null, AlignUp(Filename.size() + 1, 4) - Filename.size() - 1);
      if (!File) {
        break;
      }
      Ret[Info.ExternalFileId].Filename = std::move(Filename);
    } else if (Entry.FileId == SetExecutableFileId {}.Marker.FileId && Entry.BlockOffset == SetExecutableFileId {}.Marker.BlockOffset) {
      CodeMapFileId ExecutableFileId;
      File.read(reinterpret_cast<char*>(&ExecutableFileId), sizeof(ExecutableFileId));
      if (!File) {
        break;
      }
      Ret[ExecutableFileId].IsExecutable = true;
    } else {
      if (!Ret.contains(Entry.FileId)) {
        LogMan::Msg::EFmt("Code map referenced unknown file id {:016x}", Entry.FileId);
      } else {
        Ret[Entry.FileId].Blocks.insert(Entry.BlockOffset);
      }
    }

    if (!File) {
      break;
    }
  }
  return Ret;
}

CodeMapWriter::CodeMapWriter(CodeMapOpener& Opener, bool OpenEagerly)
  : Buffer(4096)
  , FileOpener(Opener) {
  if (OpenEagerly) {
    CodeMapFD = FileOpener.OpenCodeMapFile();
  }
}

CodeMapWriter::~CodeMapWriter() {
  if (CodeMapFD.value_or(-1) != -1) {
    Flush(BufferOffset);
    close(*CodeMapFD);
  }
}

bool CodeMapWriter::IsWriteEnabled(const ExecutableFileSectionInfo& Section) {
  if (CodeMapFD == -1) {
    return false;
  }

  // PV libraries can't yet be read by FEXServer, so skip dumping them
  if (Section.FileInfo.Filename.starts_with("/run/pressure-vessel")) {
    return false;
  }

  if (CodeMapFD) {
    return true;
  }

  // Acquire mutex and re-check CodeMapFD to avoid race conditions
  auto lk = std::unique_lock {Mutex};
  if (!CodeMapFD) {
    CodeMapFD = FileOpener.OpenCodeMapFile();
  }

  return CodeMapFD != -1;
}

void CodeMapWriter::Flush(size_t Offset) {
  // Acquire exclusive lock and flush circular buffer
  std::unique_lock Lock {Mutex};
  Flush(Offset, Lock);
}

void CodeMapWriter::Flush(size_t Offset, std::unique_lock<std::shared_mutex>&) {
  write(*CodeMapFD, Buffer.data(), Offset);
  BufferOffset = 0;
}

void CodeMapWriter::AppendBlock(const FEXCore::ExecutableFileSectionInfo& SectionInfo, uint64_t BlockEntry) {
  if (!IsWriteEnabled(SectionInfo)) {
    return;
  }

  BlockEntry -= SectionInfo.FileStartVA;
  if (BlockEntry > std::numeric_limits<uint32_t>::max()) {
    ERROR_AND_DIE_FMT("Cannot write code map");
  }

  // Register new library if not already known
  bool NewLibraryLoad = false;
  {
    // Check prior registration with shared lock
    std::shared_lock Lock {Mutex};
    NewLibraryLoad = !KnownFileIds.contains(SectionInfo.FileInfo.FileId);
  }
  if (NewLibraryLoad) {
    // Register to map with exclusive lock
    std::unique_lock Lock {Mutex};
    NewLibraryLoad &= KnownFileIds.insert(SectionInfo.FileInfo.FileId).second;
  }
  if (NewLibraryLoad) {
    // Add entry to code map
    AppendLibraryLoad(SectionInfo.FileInfo);
  }

  // Register the actual code block
  CodeMap::Entry DataEntry {SectionInfo.FileInfo.FileId, static_cast<uint32_t>(BlockEntry)};
  AppendData(std::as_bytes(std::span {&DataEntry, 1}));
}

void CodeMapWriter::AppendLibraryLoad(const FEXCore::ExecutableFileInfo& FileInfo) {
  // See CodeMap::ExternalLibraryInfo
  auto ExternalFileId = FileInfo.FileId;
  auto TotalSize = AlignUp(sizeof(CodeMap::LoadExternalLibrary) + sizeof(ExternalFileId) + FileInfo.Filename.size() + 1, 4);
  const auto Data = reinterpret_cast<char*>(alloca(TotalSize));
  auto WritePtr = std::copy_n(reinterpret_cast<const char*>(&CodeMap::LoadExternalLibrary), sizeof(CodeMap::LoadExternalLibrary), Data);
  WritePtr = std::copy_n(reinterpret_cast<const char*>(&ExternalFileId), sizeof(ExternalFileId), WritePtr);
  WritePtr = std::copy(FileInfo.Filename.begin(), FileInfo.Filename.end(), WritePtr);
  std::fill(WritePtr, Data + TotalSize, 0);
  AppendData(std::as_bytes(std::span {Data, TotalSize}));
}

void CodeMapWriter::AppendSetMainExecutable(const FEXCore::ExecutableFileInfo& FileInfo) {
  CodeMap::SetExecutableFileId Data {.ExecutableFileId = FileInfo.FileId};
  AppendData(std::span {reinterpret_cast<const std::byte*>(&Data), sizeof(Data)});
}

void CodeMapWriter::AppendData(std::span<const std::byte> Data) {
  std::shared_lock Lock {Mutex};
  auto Offset = BufferOffset.fetch_add(Data.size_bytes());
  if (Offset + Data.size_bytes() > Buffer.size()) {
    // Acquire exclusive lock and flush the buffer.
    // Under heavy pressure, multiple threads may observe an exhausted buffer simultaneously.
    // The thread with the last in-bounds Offset is responsible for flushing the buffer.
    Lock.unlock();
    bool IsResponsibleForFlush = false;
    {
      std::unique_lock ExclusiveLock {Mutex};
      IsResponsibleForFlush = (Offset <= Buffer.size());
      if (IsResponsibleForFlush) {
        Flush(Offset, ExclusiveLock);
      }
    }
    if (!IsResponsibleForFlush) {
      // Wait for the buffer to be flushed on the responsible thread
      Utils::SpinWaitLock::WaitPred<std::less_equal<>, size_t>(reinterpret_cast<size_t*>(&BufferOffset), Buffer.size());
    }
    AppendData(Data);
    return;
  }

  memcpy(&Buffer.at(Offset), Data.data(), Data.size_bytes());
}

} // namespace FEXCore

namespace FEXCore::Context {

CodeCache::CodeCache(ContextImpl& CTX_)
  : CTX(CTX_) {}
CodeCache::~CodeCache() = default;

uint64_t CodeCache::ComputeCodeMapId(std::string_view Filename, int FD) {
  if (Filename.empty()) {
    return 0xffff'ffff'ffff'ffff;
  }

  // For now, we just use the file path as an identifier.
  // TODO: Ensure the hash is unique enough to distinguish executables while remaining independent of the installation location
  return XXH3_64bits(Filename.data(), Filename.size());
}

struct CodeCacheHeader {
  std::array<char, 4> Magic = ExpectedMagic;
  // Version history:
  // 1: Initial version
  // 2: Padding code buffer data to enable direct mapping
  uint32_t FormatVersion = 2;
  uint8_t FEXVersion[20] = {};
  uint32_t NumBlocks;
  uint32_t NumCodePages;
  uint32_t CodeBufferSize;
  uint32_t NumRelocations;
  uint32_t padding;
  uint64_t SerializedBaseAddress;
  // TODO: Consider including information from LookupCache.BlockLinks

  static constexpr std::array<char, 4> ExpectedMagic = {'F', 'X', 'C', 'C'};
};

template<typename T>
concept OrderedContainer = requires { typename T::key_compare; };

bool CodeCache::SaveData(Core::InternalThreadState& Thread, int fd, const ExecutableFileSectionInfo& SourceBinary, uint64_t SerializedBaseAddress) {
  auto CodeBuffer = CTX.GetLatest();
  auto& LookupCache = *Thread.LookupCache->Shared;
  auto Relocations = Thread.CPUBackend->TakeRelocations(SourceBinary.FileStartVA);

  // Write file header
  CodeCacheHeader header {};
  static_assert(GIT_HASH.size() == sizeof(header.FEXVersion));
  std::ranges::copy(GIT_HASH, header.FEXVersion);
  header.NumBlocks = LookupCache.BlockList.size();
  header.NumCodePages = LookupCache.CodePages.size();
  header.CodeBufferSize = FEXCore::AlignUp(CTX.LatestOffset, Utils::FEX_PAGE_SIZE);
  header.NumRelocations = Relocations.size();
  header.SerializedBaseAddress = SerializedBaseAddress;
  ::write(fd, &header, sizeof(header));

  // Dump guest<->host block mappings
  {
    // Cache contents must be deterministic, so copy the unordered block list and then sort by key
    static_assert(!OrderedContainer<decltype(LookupCache.BlockList)>, "Already deterministic; drop temporary container");
    fextl::vector<std::pair<uint64_t, const GuestToHostMap::BlockEntry*>> BlockList;
    BlockList.reserve(LookupCache.BlockList.size());
    for (auto& [Guest, BlockEntry] : LookupCache.BlockList) {
      static_assert(sizeof(Guest) == 8, "Breaking change in code cache data layout");
      BlockList.emplace_back(Guest, &BlockEntry);
    }
    std::ranges::sort(BlockList);

    for (auto [Guest, Host] : BlockList) {
      static_assert(sizeof(Host->HostCode) == 8, "Breaking change in code cache data layout");
      static_assert(sizeof(Host->CodePages[0]) == 8, "Breaking change in code cache data layout");

      Guest -= SourceBinary.FileStartVA;
      ::write(fd, &Guest, sizeof(Guest));
      uint64_t HostCode = Host->HostCode - reinterpret_cast<uintptr_t>(CodeBuffer->Ptr);
      ::write(fd, &HostCode, sizeof(HostCode));
      uint64_t NumCodePages = Host->CodePages.size();
      ::write(fd, &NumCodePages, sizeof(NumCodePages));
      LOGMAN_THROW_A_FMT(std::ranges::is_sorted(Host->CodePages), "Code pages aren't sorted");
      for (auto CodePage : Host->CodePages) {
        CodePage -= SourceBinary.FileStartVA;
        ::write(fd, &CodePage, sizeof(CodePage));
      }
    }
  }

  // Dump relocations
  static_assert(sizeof(Relocations[0]) == 48, "Breaking change in code cache data layout");
  ::write(fd, Relocations.data(), Relocations.size() * sizeof(Relocations[0]));

  // Pad to next page in file so that the CodeBuffer can be mmap'ed into process on load
  {
    auto AlignedSize = AlignUp(lseek(fd, 0, SEEK_CUR), Utils::FEX_PAGE_SIZE);
    ::ftruncate(fd, AlignedSize);
    lseek(fd, AlignedSize, SEEK_SET);
  }

  // Dump the host code (relocated for position-independent serialization)
  std::span CodeBufferData(reinterpret_cast<std::byte*>(CodeBuffer->Ptr), reinterpret_cast<std::byte*>(CodeBuffer->Ptr) + CTX.LatestOffset);
  if (!ApplyCodeRelocations(SerializedBaseAddress, CodeBufferData, Relocations, 0, true)) {
    LOGMAN_THROW_A_FMT(false, "Failed to apply code relocations");
    return false;
  }
  ::write(fd, CodeBufferData.data(), CodeBufferData.size());
  // Pad to next page in file for mmap
  {
    auto PaddedSize = AlignUp(lseek(fd, 0, SEEK_CUR), Utils::FEX_PAGE_SIZE);
    ::ftruncate(fd, PaddedSize);
    lseek(fd, PaddedSize, SEEK_SET);
  }

  // Dump code pages
  static_assert(OrderedContainer<decltype(LookupCache.CodePages)>, "Non-deterministic data source");
  for (const auto& [PageIndex, Entrypoints] : LookupCache.CodePages) {
    uint64_t PageAddr = (PageIndex << 12) - SourceBinary.FileStartVA;
    ::write(fd, &PageAddr, sizeof(PageAddr));
    uint64_t NumEntrypoints = Entrypoints.size();
    ::write(fd, &NumEntrypoints, sizeof(NumEntrypoints));
    for (uint64_t Entrypoint : Entrypoints) {
      Entrypoint -= SourceBinary.FileStartVA;
      ::write(fd, &Entrypoint, sizeof(Entrypoint));
    }
  }

  return true;
}

void CodeCache::Validate(const ExecutableFileSectionInfo& Section, fextl::set<uint64_t> GuestBlocks, const fextl::set<uint64_t>& HostBlocks,
                         std::span<std::byte> CachedCode) {
  LOGMAN_THROW_A_FMT(!HostBlocks.empty(), "Tried to validate without any host blocks");
  // Skip any cached data before the first host block
  CachedCode = CachedCode.subspan(*HostBlocks.begin() - sizeof(CPU::CPUBackend::JITCodeHeader));

  if (!ValidationCTX) {
    ValidationCTX.reset(static_cast<ContextImpl*>(FEXCore::Context::Context::CreateNewContext(CTX.HostFeatures).release()));
    ValidationCTX->SetSignalDelegator(CTX.SignalDelegation);
    ValidationCTX->SetSyscallHandler(CTX.SyscallHandler);
    ValidationCTX->SetThunkHandler(CTX.ThunkHandler);
    if (!ValidationCTX->InitCore()) {
      ERROR_AND_DIE_FMT("Failed to create cache load validation context");
    }

    ValidationThread.reset(ValidationCTX->CreateThread(0, 0, nullptr));

    auto Frame = ValidationThread->CurrentFrame;
    Frame->State.segment_arrays[FEXCore::Core::CPUState::SEGMENT_ARRAY_INDEX_GDT] = &ValidationGDT[0];
    Frame->State.segment_arrays[FEXCore::Core::CPUState::SEGMENT_ARRAY_INDEX_LDT] = &ValidationGDT[0];
    Frame->State.cs_idx = 0;
    Frame->State.cs_cached = 0;

    if (ValidationCTX->Config.Is64BitMode()) {
      ValidationGDT[0].L = 1; // L = Long Mode = 64-bit
      ValidationGDT[0].D = 0; // D = Default Operand Size = Reserved
    } else {
      ValidationGDT[0].L = 0; // L = Long Mode = 32-bit
      ValidationGDT[0].D = 1; // D = Default Operand Size = 32-bit
    }
  }

  auto NewCodeBuffer = ValidationCTX->GetLatest();
  while (CachedCode.size_bytes() > NewCodeBuffer->UsableSize()) {
    ValidationCTX->ClearCodeCache(ValidationThread.get());
    NewCodeBuffer = ValidationCTX->GetLatest();
    LogMan::Msg::IFmt("Increased cache validation code buffer size to {} MiB", NewCodeBuffer->AllocatedSize / 1024 / 1024);
  }

  std::span<std::byte> CodeBufferRangeRef =
    std::as_writable_bytes(std::span {NewCodeBuffer->Ptr, NewCodeBuffer->Ptr + NewCodeBuffer->UsableSize()}).subspan(0, CachedCode.size_bytes());

  while (!GuestBlocks.empty()) {
    auto [CompiledBlocks, _, _2, _3, _4] = ValidationCTX->CompileCode(ValidationThread.get(), *GuestBlocks.begin(), 0 /* TODO: Set MaxInst? */);
    for (auto& Entry : CompiledBlocks.EntryPoints) {
      GuestBlocks.erase(Entry.first);
    }
  }

  // Patch FEX-internal function addresses with values from the main Context to ensure the code blocks are comparable
  auto NewRelocations = ValidationThread->CPUBackend->TakeRelocations(Section.FileStartVA);
  NewRelocations.erase(std::remove_if(NewRelocations.begin(), NewRelocations.end(), [](const CPU::Relocation& Reloc) {
    return Reloc.Header.Type != CPU::RelocationTypes::RELOC_NAMED_SYMBOL_LITERAL && Reloc.Header.Type != CPU::RelocationTypes::RELOC_NAMED_THUNK_MOVE;
  }));
  (void)ApplyCodeRelocations(Section.FileStartVA, CodeBufferRangeRef, NewRelocations, 0, false);

  if (ValidationCTX->LatestOffset <= CodeBufferRangeRef.size()) {
    // Reference compilation produced fewer bytes than our cache, so validation is going to fail.
    // Make sure we don't output any garbage bytes though.
    CodeBufferRangeRef = CodeBufferRangeRef.subspan(0, ValidationCTX->LatestOffset);
  }

  auto [Mismatch, _] = std::mismatch(CodeBufferRangeRef.begin(), CodeBufferRangeRef.end(), CachedCode.begin());
  if (Mismatch != CodeBufferRangeRef.end()) {
    // Align down to instruction size
    auto Idx = AlignDown(std::distance(CodeBufferRangeRef.begin(), Mismatch), 4);

    auto BlockIt = std::prev(HostBlocks.lower_bound(*HostBlocks.begin() + Idx + 1));
    std::optional<uint64_t> GuestBlockAddr;
    std::optional<uint64_t> GuestBlockAddrRef;
    if (BlockIt != HostBlocks.end()) {
      for (int i : {0, 1}) {
        std::span Buffer = (i == 0 ? CachedCode : CodeBufferRangeRef);

        // Second instruction is always a constant load for relative offset to the (multi)block start
        int32_t addr = (*reinterpret_cast<uint32_t*>(&Buffer[*BlockIt - *HostBlocks.begin() + 4]) & 0x3ff'ffe0) << 11;
        addr >>= 14;
        auto header = reinterpret_cast<CPU::CPUBackend::JITCodeHeader*>(&Buffer[*BlockIt - *HostBlocks.begin() + 4 + addr]);
        auto tail = reinterpret_cast<CPU::CPUBackend::JITCodeTail*>(reinterpret_cast<uintptr_t>(header) + header->OffsetToBlockTail);
        (i == 0 ? GuestBlockAddr : GuestBlockAddrRef) = tail->RIP - Section.FileStartVA;
        LogMan::Msg::EFmt("Recorded rip {}: {:#x} (offset {:#x})", i, tail->RIP, tail->RIP - Section.FileStartVA);

        if (i == 1) {
          if (tail->RIP >= Section.BeginVA && tail->RIP < Section.EndVA) {
            auto [IRView, TotalInstructions, TotalInstructionsLength, StartAddr, Length, _] =
              ValidationCTX->GenerateIR(ValidationThread.get(), tail->RIP, false, FEXCore::Config::Get_MAXINST());
            fextl::stringstream ss;
            FEXCore::IR::Dump(&ss, &*IRView);
            LogMan::Msg::EFmt("IR:\n{}", ss.str());
          } else {
            LogMan::Msg::EFmt("Can't dump IR for out-of-range RIP {:#x}", tail->RIP);
          }
        }
      }
    }

    fextl::string GuestBlockInfo = "UNKNOWN";
    if (GuestBlockAddr) {
      GuestBlockInfo = fextl::fmt::format("{:#x}", GuestBlockAddr.value());
    }
    if (GuestBlockAddr != GuestBlockAddrRef) {
      GuestBlockInfo += " (MISMATCH)";
    }
    ERROR_AND_DIE_FMT("Cache validation failed at offset {:#x}: {:02x} <-> {:02x} (at {} <-> {}, guest block {})", Idx,
                      fmt::join(CachedCode.subspan(Idx, 4), ""), fmt::join(CodeBufferRangeRef.subspan(Idx, 4), ""),
                      fmt::ptr(CachedCode.data()), fmt::ptr(CodeBufferRangeRef.data()), GuestBlockInfo);
  }

  // Reset Context state for next validation
  ValidationThread->LookupCache->ClearCache(ValidationThread->LookupCache->AcquireWriteLock());
  ValidationCTX->LatestOffset = 0;

  LogMan::Msg::IFmt("            successfully validated cache");
}

bool CodeCache::ApplyCodeRelocations(uint64_t GuestEntry, std::span<std::byte> Code,
                                     std::span<const FEXCore::CPU::Relocation> EntryRelocations, uint32_t RelocationOffset, bool ForStorage) {
  CPU::Arm64Emitter Emitter(&CTX, Code.data(), Code.size_bytes());
  for (size_t j = 0; j < EntryRelocations.size(); ++j) {
    const FEXCore::CPU::Relocation& Reloc = EntryRelocations[j];
    LOGMAN_THROW_A_FMT(Reloc.Header.Offset >= RelocationOffset, "Invalid relocation offset");
    LOGMAN_THROW_A_FMT(Reloc.Header.Offset - RelocationOffset < Code.size_bytes(), "Invalid relocation offset");
    Emitter.SetCursorOffset(Reloc.Header.Offset - RelocationOffset);

    switch (Reloc.Header.Type) {
    case FEXCore::CPU::RelocationTypes::RELOC_NAMED_SYMBOL_LITERAL: {
      // Generate a literal so we can place it
      uint64_t Pointer = ForStorage ? 0 : GetNamedSymbolLiteral(CTX, Reloc.NamedSymbolLiteral.Symbol);
      Emitter.dc64(Pointer);
      break;
    }
    case FEXCore::CPU::RelocationTypes::RELOC_NAMED_THUNK_MOVE: {
      uint64_t Pointer = ForStorage ? 0 : reinterpret_cast<uint64_t>(CTX.ThunkHandler->LookupThunk(Reloc.NamedThunkMove.Symbol));
      if (Pointer == ~0ULL) {
        return false;
      }
      // TODO: Pointers are required to fit within 48-bit VA space.
      // But forcing 6-byte broke relocations.
      Emitter.LoadConstant(ARMEmitter::Size::i64Bit, ARMEmitter::Register(Reloc.NamedThunkMove.RegisterIndex), Pointer,
                           CPU::Arm64Emitter::PadType::DOPAD);
      break;
    }
    case FEXCore::CPU::RelocationTypes::RELOC_GUEST_RIP_LITERAL: {
      Emitter.dc64(GuestEntry + Reloc.GuestRIP.GuestRIP);
      break;
    }
    case FEXCore::CPU::RelocationTypes::RELOC_GUEST_RIP_MOVE: {
      uint64_t Pointer = Reloc.GuestRIP.GuestRIP + GuestEntry;
      // TODO: Pointers are required to fit within 48-bit VA space.
      // But forcing 6-byte broke relocations.
      Emitter.LoadConstant(ARMEmitter::Size::i64Bit, ARMEmitter::Register(Reloc.GuestRIP.RegisterIndex), Pointer, CPU::Arm64Emitter::PadType::DOPAD);
      break;
    }

    default: ERROR_AND_DIE_FMT("Unknown relocation type {}", ToUnderlying(Reloc.Header.Type));
    }
  }

  return true;
}

fextl::unique_ptr<MappedCodeCacheFile>
CodeCache::LoadCache(std::span<std::byte> CacheFile, const ExecutableFileInfo& FileInfo, uint64_t FileStartVA) {
  if (!EnableCodeCaching) {
    return nullptr;
  }

  FEXCORE_PROFILE_SCOPED("LoadCache");

  // Read file header
  CodeCacheHeader header {};
  ::memcpy(&header, CacheFile.data(), sizeof(header));

  if (!std::ranges::equal(header.Magic, header.ExpectedMagic)) {
    LogMan::Msg::EFmt("Invalid cache file header");
    return nullptr;
  }

  if (!std::ranges::equal(header.FEXVersion, GIT_HASH)) {
    LogMan::Msg::IFmt("Cache generated from old FEX version {:02x}, current is {:02x}; skipping", fmt::join(header.FEXVersion, ""),
                      fmt::join(GIT_HASH, ""));
    return nullptr;
  }

  if (header.NumBlocks == 0) {
    // Valid caches are never empty
    LogMan::Msg::IFmt("Code cache empty, aborting");
    return nullptr;
  }

  // Skip over BlockEntry data since it won't be used until EnableLoadedSection
  // TODO: Store direct offset to relocations in the header
  auto* BlockListStart = CacheFile.data() + sizeof(header);
  auto* Cursor = BlockListStart;
  for (uint32_t i = 0; i < header.NumBlocks; ++i) {
    Cursor += sizeof(uint64_t); // guest address
    Cursor += sizeof(uint64_t); // host code address
    uint64_t NumGuestCodePages;
    ::memcpy(&NumGuestCodePages, Cursor, sizeof(NumGuestCodePages));
    Cursor += sizeof(NumGuestCodePages);
    Cursor += NumGuestCodePages * sizeof(uint64_t);
  }

  auto Relocations = std::span {reinterpret_cast<const FEXCore::CPU::Relocation*>(Cursor), header.NumRelocations};
  Cursor += Relocations.size_bytes();

  // Pad to next page to get the code buffer data
  Cursor = reinterpret_cast<std::byte*>(AlignUp(reinterpret_cast<uintptr_t>(Cursor), Utils::FEX_PAGE_SIZE));
  auto CodeDataInFile = std::span {Cursor, header.CodeBufferSize};

#ifndef _WIN32
  // Allocate target memory for post-relocation code. This is PROT_NONE until
  // the first execution, so that contents can be lazily populated in a
  // frontend-provided segfault handler.
  void* CodeBufferAllocation = Allocator::mmap(nullptr, header.CodeBufferSize, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (CodeBufferAllocation == MAP_FAILED) {
    LogMan::Msg::EFmt("Failed to reserve target memory for code cache");
    return nullptr;
  }
  auto CodeBuffer = std::span {static_cast<std::byte*>(CodeBufferAllocation), header.CodeBufferSize};
#else
  // TODO: Implement lazy mapping on Windows
  auto CodeBuffer = CodeDataInFile;
#endif

  // Group relocations by page
  size_t NumPages = header.CodeBufferSize / Utils::FEX_PAGE_SIZE;
  fextl::vector<MappedCodeCacheFile::PageRelocationRange> PageRelocationRanges(NumPages, {0, 0});
  auto RelocBaseOffset = std::as_bytes(Relocations).data() - CacheFile.data();
  auto RelocIt = Relocations.begin();
  for (size_t Page = 0; Page < NumPages; ++Page) {
    auto EndRelocIt = std::upper_bound(RelocIt, Relocations.end(), Page,
                                       [](auto& Page, auto& Reloc) { return Page < Reloc.Header.Offset / Utils::FEX_PAGE_SIZE; });
    PageRelocationRanges.at(Page) = {static_cast<uint32_t>(RelocBaseOffset + (RelocIt - Relocations.begin()) * sizeof(CPU::Relocation)),
                                     static_cast<uint32_t>(EndRelocIt - RelocIt)};
    RelocIt = EndRelocIt;
  }

  auto Storage = FEXCore::Allocator::aligned_alloc(alignof(MappedCodeCacheFile), sizeof(MappedCodeCacheFile));
  return fextl::unique_ptr<MappedCodeCacheFile>(
    new (Storage) MappedCodeCacheFile {this, CacheFile, CodeDataInFile, CodeBuffer, BlockListStart, header.NumBlocks, header.NumCodePages,
                                       std::move(PageRelocationRanges), fextl::vector<bool>(NumPages), FileStartVA});
}

bool CodeCache::EnableLoadedSection(Core::InternalThreadState* Thread, MappedCodeCacheFile& Code, const ExecutableFileSectionInfo& BinarySection) {
  if (!EnableCodeCaching) {
    return true;
  }

  namespace ranges = std::ranges;

  FEXCORE_PROFILE_SCOPED("EnableLoadedSection");

  // Read block list from cache file
  // TODO: Store section-ized BlockLists in cache file
  using BlockListEntry = decltype(GuestToHostMap::BlockList)::value_type;
  fextl::vector<BlockListEntry> BlockList(Code.NumBlocks);
  {
    auto* Cursor = Code.BlockListInFile;
    for (auto& BlockPtr : BlockList) {
      ::memcpy(&BlockPtr.first, Cursor, sizeof(BlockPtr.first));
      Cursor += sizeof(BlockPtr.first);
      ::memcpy(&BlockPtr.second.HostCode, Cursor, sizeof(BlockPtr.second.HostCode));
      Cursor += sizeof(BlockPtr.second.HostCode);
      uint64_t NumGuestPages;
      ::memcpy(&NumGuestPages, Cursor, sizeof(NumGuestPages));
      Cursor += sizeof(NumGuestPages);

      BlockPtr.second.CodePages.resize(NumGuestPages);
      ::memcpy(BlockPtr.second.CodePages.data(), Cursor, std::span {BlockPtr.second.CodePages}.size_bytes());
      Cursor += std::span {BlockPtr.second.CodePages}.size_bytes();
    }

    // Constrain BlockList to the given ExecutableFileSectionInfo
    LOGMAN_THROW_A_FMT(ranges::is_sorted(BlockList, [](auto& a, auto& b) { return a.first < b.first; }), "Expected sorted block list");
    auto begin = ranges::lower_bound(BlockList, BinarySection.BeginVA - BinarySection.FileStartVA, std::less {}, &BlockListEntry::first);
    auto end =
      ranges::upper_bound(begin, BlockList.end(), BinarySection.EndVA - BinarySection.FileStartVA - 1, std::less {}, &BlockListEntry::first);
    if (begin == end) {
      LogMan::Msg::IFmt("No blocks cached in this range, aborting");
      return true;
    }
    BlockList.erase(end, BlockList.end());
    BlockList.erase(BlockList.begin(), begin);
  }

  LogMan::Msg::IFmt("Cache load: {:5} blocks; base={:#14x}; off={:#9x}-{:#09x}; {:016x} {}", BlockList.size(), BinarySection.FileStartVA,
                    BinarySection.BeginVA - BinarySection.FileStartVA, BinarySection.EndVA - BinarySection.FileStartVA,
                    BinarySection.FileInfo.FileId, BinarySection.FileInfo.Filename);

  if (EnableLazyCodeCaching) {
    LogMan::Msg::IFmt("            lazy mapping: base={:#14x} -> host={}; cache_source={}", BinarySection.FileStartVA,
                      fmt::ptr(Code.CodeBuffer.data()), fmt::ptr(Code.MappedFile.data()));
  }
  // Register blocks to LookupCache.
  // The host addresses will point into the protected code buffer, so that FEX
  // can lazily apply relocations on first execution of each page.
  auto CodeBuffer = CTX.GetLatest();
  {
    FEXCORE_PROFILE_SCOPED("Decode");
    auto& LookupCache = *CodeBuffer->LookupCache;
    auto WriteLock = LookupCache.AcquireWriteLock();

    for (auto& [Guest, Block] : BlockList) {
      for (auto& CodePage : Block.CodePages) {
        CodePage += BinarySection.FileStartVA;
      }
      LOGMAN_THROW_A_FMT(Block.HostCode < Code.CodeBuffer.size_bytes(), "Host offset {:#x} out of range ({:#x})", Block.HostCode,
                         Code.CodeBuffer.size_bytes());
      auto HostCode = &Code.CodeBuffer[Block.HostCode];
      LookupCache.AddBlockMapping(Guest + BinarySection.FileStartVA, std::move(Block.CodePages), HostCode, WriteLock);
    }

    // Guest code pages
    auto* Cursor = Code.CodeBufferInFile.data() + Code.CodeBufferInFile.size_bytes();
    fextl::vector<uint64_t> Entrypoints;
    for (uint32_t i = 0; i < Code.NumCodePages; ++i) {
      uint64_t CodePage;
      memcpy(&CodePage, Cursor, sizeof(CodePage));
      CodePage += BinarySection.FileStartVA;
      Cursor += sizeof(CodePage);

      uint64_t NumEntrypoints;
      memcpy(&NumEntrypoints, Cursor, sizeof(NumEntrypoints));
      Cursor += sizeof(NumEntrypoints);

      Entrypoints.resize(NumEntrypoints);
      memcpy(Entrypoints.data(), Cursor, std::span {Entrypoints}.size_bytes());
      Cursor += std::span {Entrypoints}.size_bytes();
      for (auto& Entrypoint : Entrypoints) {
        Entrypoint += BinarySection.FileStartVA;
      }

      if (LookupCache.AddBlockExecutableRange(Entrypoints, CodePage, FEXCore::Utils::FEX_PAGE_SIZE, WriteLock)) {
        CTX.SyscallHandler->MarkGuestExecutableRange(Thread, CodePage, FEXCore::Utils::FEX_PAGE_SIZE);
      }
    }
  }

#ifndef _WIN32
  if (!EnableLazyCodeCaching || EnableCodeCacheValidation) {
#else
  // TODO: Implement lazy mapping on Windows
  if (true) {
#endif
    auto Range = SelectCodeRangeToFinalize(Code, 0, Code.CodeBuffer.size_bytes() / Utils::FEX_PAGE_SIZE);
    FinalizeCodePages(Code, Range);
  }

  if (EnableCodeCacheValidation) {
    fextl::set<uint64_t> GuestBlocks, HostBlocks;
    for (auto& [Guest, Host] : BlockList) {
      GuestBlocks.insert(Guest + BinarySection.FileStartVA);
      HostBlocks.insert(Host.HostCode);
    }

    Validate(BinarySection, std::move(GuestBlocks), HostBlocks, Code.CodeBuffer);
  }

  return true;
}

} // namespace FEXCore::Context

namespace FEXCore {

static std::span<CPU::Relocation> SpanPageRelocations(const MappedCodeCacheFile& Code, size_t PageIndex) {
  auto [Offset, Count] = Code.PageRelocationRanges.at(PageIndex);
  return std::span {reinterpret_cast<FEXCore::CPU::Relocation*>(Code.MappedFile.data() + Offset), Count};
}

std::span<std::byte> AbstractCodeCache::SelectCodeRangeToFinalize(MappedCodeCacheFile& Code, size_t StartPage, size_t EndPage) {
  // First, check if we were racing another thread in loading this range
  if (std::find(Code.LoadedPages.begin() + StartPage, Code.LoadedPages.begin() + EndPage, false) == Code.LoadedPages.begin() + EndPage) {
    return {};
  }

  LOGMAN_THROW_A_FMT(StartPage < EndPage, "Invalid page range [{}, {})", StartPage, EndPage);
  LOGMAN_THROW_A_FMT(EndPage <= Code.NumPages(), "End page {} out of range ({})", EndPage, Code.NumPages());

  // Include any pages that have relocations or block link records crossing
  // into the current page range. This ensures we don't attempt to finalize
  // any page twice, partially apply FEX relocations, or trigger page loads
  // during block linking.
  while (EndPage < Code.NumPages()) {
    auto PageRelocs = SpanPageRelocations(Code, EndPage - 1);
    if (!PageRelocs.empty()) {
      auto It = std::prev(PageRelocs.end());
      size_t RelocEnd = It->Header.Offset + 16 /* Upper bound for relocation size */;
      if (RelocEnd > EndPage * Utils::FEX_PAGE_SIZE) {
        ++EndPage;
        continue;
      }
    }

    // Check for trailing block link
    {
      auto PageRelocs = SpanPageRelocations(Code, EndPage);
      if (!PageRelocs.empty() && PageRelocs.begin()->Header.Offset < EndPage * Utils::FEX_PAGE_SIZE + 0x18) {
        ++EndPage;
        continue;
      }
    }
    break;
  };
  while (StartPage != 0) {
    auto PageRelocs = SpanPageRelocations(Code, StartPage - 1);
    if (!PageRelocs.empty()) {
      auto It = std::prev(PageRelocs.end());
      size_t RelocEnd = It->Header.Offset + 16 /* Upper bound for relocation size */;
      if (RelocEnd > StartPage * Utils::FEX_PAGE_SIZE) {
        --StartPage;
        continue;
      }
    }

    // Check for trailing block link
    {
      auto PageRelocs = SpanPageRelocations(Code, StartPage);
      if (!PageRelocs.empty() && PageRelocs.begin()->Header.Offset < StartPage * Utils::FEX_PAGE_SIZE + 0x18) {
        --StartPage;
        continue;
      }
    }
    break;
  };

  return Code.CodeBuffer.subspan(StartPage * Utils::FEX_PAGE_SIZE, (EndPage - StartPage) * Utils::FEX_PAGE_SIZE);
}
} // namespace FEXCore

namespace FEXCore::Context {

void CodeCache::FinalizeCodePages(MappedCodeCacheFile& Code, std::span<std::byte> CodeRange) {
  const size_t StartOffset = CodeRange.data() - Code.CodeBuffer.data();
  const auto StartPage = StartOffset / Utils::FEX_PAGE_SIZE;
  const auto EndPage = StartPage + CodeRange.size_bytes() / Utils::FEX_PAGE_SIZE;
  const size_t Size = CodeRange.size_bytes();

  // None of the selected pages should be loaded at all; otherwise, SelectCodeRangeToFinalize returned inconsistent ranges
  LOGMAN_THROW_A_FMT(std::find(Code.LoadedPages.begin() + StartPage, Code.LoadedPages.begin() + EndPage, true) == Code.LoadedPages.begin() + EndPage,
                     "Inconsistent page load state");

  FEXCORE_PROFILE_SCOPED("FinalizeCodePages");

#ifndef _WIN32
  // Atomicity is critical when making the finalized code data visible.
  // We ensure this by remapping a temporary buffer onto the PROT_NONE
  // placeholder page in CodeBuffer. Some constraints to keep in mind are:
  // 1. Pages can't be write-only (readability is implicitly added), so
  //    we can't change CodeBuffer from PROT_NONE to PROT_WRITE even for just
  //    a short duration
  // 2. Naive mremap from CodeBufferInFile to CodeBuffer would leave a gap in
  //    the former, which would make cleanup overly complicated
  //
  // Due to (1), we can't apply relocations in place (CodeBufferInFile); at
  // least a secondary buffer is needed for execution (CodeBuffer).
  // Due to (2), a third buffer is temporarily allocated here and freed on
  // completion. The final code data is computed here and then the memory
  // is remapped onto CodeBuffer.
  auto* Staging = reinterpret_cast<std::byte*>(Allocator::VirtualAlloc(nullptr, Size, true));
  if (!Staging) {
    ERROR_AND_DIE_FMT("Failed to allocate {} bytes of staging memory for code-cache finalization", Size);
  }

  // Copy code from the cache file to the staging buffer
  memcpy(Staging, Code.CodeBufferInFile.data() + StartOffset, Size);

  // Apply relocations
  auto StagingSpan = std::span {Staging, Size};
  for (size_t i = StartPage; i < EndPage; ++i) {
    auto PageRelocations = SpanPageRelocations(Code, i);
    (void)ApplyCodeRelocations(Code.GuestBase, StagingSpan, PageRelocations, static_cast<uint32_t>(StartOffset), false);
    Code.LoadedPages[i] = true;
  }

  // Atomically make the finalized code data visible by remapping the staging
  // buffer onto the requested CodeBuffer window. MREMAP_DONTUNMAP is used to
  // leave the old VA range reserved so that we can cleanly deallocate it
  // through Allocator.
  void* RemapResult = ::mremap(Staging, Size, Size, MREMAP_FIXED | MREMAP_MAYMOVE | MREMAP_DONTUNMAP, CodeRange.data());
  if (RemapResult == MAP_FAILED) {
    ERROR_AND_DIE_FMT("{}: mremap failed: {}", __FUNCTION__, errno);
  }
  Allocator::VirtualFree(Staging, Size);

  // Release resident file pages that will no longer be needed. The VA range is left allocated to allow cleanup with a single VirtualFree.
  Allocator::VirtualDontNeed(Code.CodeBufferInFile.data() + StartOffset, Size);
#else
  // TODO: Implement lazy mapping on Windows
  for (size_t i = StartPage; i < EndPage; ++i) {
    auto PageRelocations = SpanPageRelocations(Code, i);
    (void)ApplyCodeRelocations(Code.GuestBase, Code.CodeBuffer, PageRelocations, 0, false);
    Code.LoadedPages[i] = true;
  }
#endif

  ARMEmitter::Emitter::ClearICache(CodeRange.data(), Size);
}

} // namespace FEXCore::Context