From f964a5187e9b0c9af2442c4258eb9c2bfc9f5a48 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Thu, 11 Jul 2024 13:40:38 +0100
Subject: [PATCH] ARM64EC: Implement BeginSimulation

This is used by the kernel (or UNIX side of ntdll in wine) to jump into
x86 code with the given context as is necessary when e.g. returning from
an exception.
---
 Source/Windows/ARM64EC/Module.S   |  9 +++++
 Source/Windows/ARM64EC/Module.cpp | 59 +++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/Source/Windows/ARM64EC/Module.S b/Source/Windows/ARM64EC/Module.S
index b7d8516fa..1ed2bd599 100644
--- a/Source/Windows/ARM64EC/Module.S
+++ b/Source/Windows/ARM64EC/Module.S
@@ -37,6 +37,15 @@ enter_jit:
   ldr x16, [x17, #0x40] // ChpeV2CpuAreaInfo->EmulatorData[2] - DispatcherLoopTopEnterEC
   br x16 // DispatcherLoopTopEnterEC(RIP:x9, CPUArea:x17)
 
+  // Invoked by KiUserEmulationDispatcher after e.g. an NtContinue to x86 code
+  // Expects a CONTEXT pointer in x0
+.global BeginSimulation
+BeginSimulation:
+  bl "#SyncThreadContext"
+  ldr x17, [x18, #0x1788] // TEB->ChpeV2CpuAreaInfo
+  ldr x16, [x17, #0x48] // ChpeV2CpuAreaInfo->EmulatorData[3] - DispatcherLoopTopEnterECFillSRA
+  br x16 // DispatcherLoopTopEnterECFillSRA(CPUArea:x17)
+
   // Called into by FEXCore
   // Expects the target code address in x9
 .global ExitFunctionEC
diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp
index 7e5c22673..02e8d08bd 100644
--- a/Source/Windows/ARM64EC/Module.cpp
+++ b/Source/Windows/ARM64EC/Module.cpp
@@ -153,6 +153,60 @@ static bool HandleUnalignedAccess(ARM64_NT_CONTEXT& Context) {
   Context.Pc += Result.second;
   return true;
 }
+
+static void LoadStateFromECContext(FEXCore::Core::InternalThreadState* Thread, CONTEXT& Context) {
+  auto& State = Thread->CurrentFrame->State;
+
+  // General register state
+  State.gregs[FEXCore::X86State::REG_RAX] = Context.Rax;
+  State.gregs[FEXCore::X86State::REG_RCX] = Context.Rcx;
+  State.gregs[FEXCore::X86State::REG_RDX] = Context.Rdx;
+  State.gregs[FEXCore::X86State::REG_RBX] = Context.Rbx;
+  State.gregs[FEXCore::X86State::REG_RSP] = Context.Rsp;
+  State.gregs[FEXCore::X86State::REG_RBP] = Context.Rbp;
+  State.gregs[FEXCore::X86State::REG_RSI] = Context.Rsi;
+  State.gregs[FEXCore::X86State::REG_RDI] = Context.Rdi;
+  State.gregs[FEXCore::X86State::REG_R8] = Context.R8;
+  State.gregs[FEXCore::X86State::REG_R9] = Context.R9;
+  State.gregs[FEXCore::X86State::REG_R10] = Context.R10;
+  State.gregs[FEXCore::X86State::REG_R11] = Context.R11;
+  State.gregs[FEXCore::X86State::REG_R12] = Context.R12;
+  State.gregs[FEXCore::X86State::REG_R13] = Context.R13;
+  State.gregs[FEXCore::X86State::REG_R14] = Context.R14;
+  State.gregs[FEXCore::X86State::REG_R15] = Context.R15;
+
+  State.rip = Context.Rip;
+  CTX->SetFlagsFromCompactedEFLAGS(Thread, Context.EFlags);
+
+  State.es_idx = Context.SegEs & 0xffff;
+  State.cs_idx = Context.SegCs & 0xffff;
+  State.ss_idx = Context.SegSs & 0xffff;
+  State.ds_idx = Context.SegDs & 0xffff;
+  State.fs_idx = Context.SegFs & 0xffff;
+  State.gs_idx = Context.SegGs & 0xffff;
+
+  // The TEB is the only populated GDT entry by default
+  const auto TEB = reinterpret_cast<uint64_t>(NtCurrentTeb());
+  State.gdt[(Context.SegGs & 0xffff) >> 3].base = TEB;
+  State.gs_cached = TEB;
+  State.fs_cached = 0;
+  State.es_cached = 0;
+  State.cs_cached = 0;
+  State.ss_cached = 0;
+  State.ds_cached = 0;
+
+  // Floating-point register state
+  CTX->SetXMMRegistersFromState(Thread, reinterpret_cast<const __uint128_t*>(Context.FltSave.XmmRegisters), nullptr);
+  memcpy(State.mm, Context.FltSave.FloatRegisters, sizeof(State.mm));
+
+  State.FCW = Context.FltSave.ControlWord;
+  State.flags[FEXCore::X86State::X87FLAG_C0_LOC] = (Context.FltSave.StatusWord >> 8) & 1;
+  State.flags[FEXCore::X86State::X87FLAG_C1_LOC] = (Context.FltSave.StatusWord >> 9) & 1;
+  State.flags[FEXCore::X86State::X87FLAG_C2_LOC] = (Context.FltSave.StatusWord >> 10) & 1;
+  State.flags[FEXCore::X86State::X87FLAG_C3_LOC] = (Context.FltSave.StatusWord >> 14) & 1;
+  State.flags[FEXCore::X86State::X87FLAG_TOP_LOC] = (Context.FltSave.StatusWord >> 11) & 0b111;
+  State.AbridgedFTW = Context.FltSave.TagWord;
+}
 } // namespace Exception
 
 namespace Logging {
@@ -195,6 +249,11 @@ public:
   }
 };
 
+extern "C" void SyncThreadContext(CONTEXT* Context) {
+  auto* Thread = GetCPUArea().ThreadState();
+  Exception::LoadStateFromECContext(Thread, *Context);
+}
+
 void ProcessInit() {
   Logging::Init();
   FEX::Config::InitializeConfigs();