From 4986b03007cf1ec255c14061c8a50cedabc2d97e Mon Sep 17 00:00:00 2001
From: Neil Dhar <neildhar@meta.com>
Date: Wed, 13 Nov 2024 11:18:06 -0800
Subject: [PATCH] Populate all outgoing call registers in JIT emitted code

Summary:
Populate the registers in the emitted code instead of the handler, so
that we can make calls to JIT functions directly in a later diff.

Reviewed By: tmikov

Differential Revision: D64275850

fbshipit-source-id: 3d4505863e88e8e8355a46f01c3a64ffe1eb4f85
---
 lib/VM/JIT/arm64/JitEmitter.cpp  | 37 ++++++++++++++++++++++++++++++--
 lib/VM/JIT/arm64/JitHandlers.cpp | 14 ++++--------
 lib/VM/JIT/arm64/JitHandlers.h   |  5 ++---
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/lib/VM/JIT/arm64/JitEmitter.cpp b/lib/VM/JIT/arm64/JitEmitter.cpp
index e9df15ba9ec..2c554b60129 100644
--- a/lib/VM/JIT/arm64/JitEmitter.cpp
+++ b/lib/VM/JIT/arm64/JitEmitter.cpp
@@ -3629,6 +3629,37 @@ void Emitter::callImpl(FR frRes, FR frCallee) {
         calleeFrameArg, calleeReg, frameRegs_[frCallee.index()].localType);
   }
 
+  static_assert(
+      HERMESVALUE_VERSION == 1,
+      "Native pointers must be encoded without modification");
+
+  FR previousFrameArg{nRegs + hbc::StackFrameLayout::PreviousFrame};
+  // Free any existing temp so we store directly.
+  freeFRTemp(previousFrameArg);
+  movFRFromHW(previousFrameArg, HWReg(xFrame), FRType::OtherNonPtr);
+
+  FR savedIPArg{nRegs + hbc::StackFrameLayout::SavedIP};
+  // Since we need a register to compute the IP in anyway, it is convenient to
+  // just use any existing one for the SavedIP slot, and let the syncAllFRTemp
+  // below write it to memory.
+  auto savedIPReg = getOrAllocFRInGpX(savedIPArg, false);
+
+  // Save the current IP in both the SavedIP slot and the runtime.
+  getBytecodeIP(savedIPReg.a64GpX());
+  frUpdatedWithHW(savedIPArg, savedIPReg, FRType::OtherNonPtr);
+  a.str(savedIPReg.a64GpX(), a64::Mem(xRuntime, offsetof(Runtime, currentIP_)));
+
+  FR savedCodeBlockArg = FR{nRegs + hbc::StackFrameLayout::SavedCodeBlock};
+  // TODO: We should be able to directly store xzr.
+  auto savedCodeBlockReg = getOrAllocFRInGpX(savedCodeBlockArg, false);
+  frUpdatedWithHW(savedCodeBlockArg, savedCodeBlockReg, FRType::OtherNonPtr);
+  a.mov(savedCodeBlockReg.a64GpX(), 0);
+
+  FR shLocalsArg{nRegs + hbc::StackFrameLayout::SHLocals};
+  // Free any existing temp so we store directly.
+  freeFRTemp(shLocalsArg);
+  movFRFromHW(shLocalsArg, savedCodeBlockReg, FRType::OtherNonPtr);
+
 #ifndef NDEBUG
   // No need to sync the set up call stack to the frame memory,
   // because it these registers can't have global registers.
@@ -3643,8 +3674,10 @@ void Emitter::callImpl(FR frRes, FR frCallee) {
   freeAllFRTempExcept({});
 
   a.mov(a64::x0, xRuntime);
-  a.mov(a64::x1, xFrame);
-  EMIT_RUNTIME_CALL(
+  loadFrameAddr(
+      a64::x1,
+      FR(frameRegs_.size() + hbc::StackFrameLayout::CalleeClosureOrCB));
+  EMIT_RUNTIME_CALL_WITHOUT_SAVED_IP(
       *this,
       SHLegacyValue(*)(SHRuntime *, SHLegacyValue *),
       _jit_dispatch_call);
diff --git a/lib/VM/JIT/arm64/JitHandlers.cpp b/lib/VM/JIT/arm64/JitHandlers.cpp
index c3e7d98f2c1..cc313d9e8cc 100644
--- a/lib/VM/JIT/arm64/JitHandlers.cpp
+++ b/lib/VM/JIT/arm64/JitHandlers.cpp
@@ -264,18 +264,12 @@ void *_jit_find_catch_target(
   _sh_throw_current(shr);
 }
 
-SHLegacyValue _jit_dispatch_call(SHRuntime *shr, SHLegacyValue *frame) {
+SHLegacyValue _jit_dispatch_call(
+    SHRuntime *shr,
+    SHLegacyValue *callTargetSHLV) {
   Runtime &runtime = getRuntime(shr);
 
-  // TODO: Move this call setup and the fast path into the emitted JIT code.
-  StackFramePtr newFrame(runtime.getStackPointer());
-  newFrame.getPreviousFrameRef() = HermesValue::encodeNativePointer(frame);
-  newFrame.getSavedIPRef() =
-      HermesValue::encodeNativePointer(runtime.getCurrentIP());
-  newFrame.getSavedCodeBlockRef() = HermesValue::encodeNativePointer(nullptr);
-  newFrame.getSHLocalsRef() = HermesValue::encodeNativePointer(nullptr);
-
-  auto *callTarget = &newFrame.getCalleeClosureOrCBRef();
+  auto *callTarget = toPHV(callTargetSHLV);
   if (vmisa<JSFunction>(*callTarget)) {
     JSFunction *jsFunc = vmcast<JSFunction>(*callTarget);
     if (auto *fnPtr = jsFunc->getCodeBlock()->getJITCompiled())
diff --git a/lib/VM/JIT/arm64/JitHandlers.h b/lib/VM/JIT/arm64/JitHandlers.h
index 99778e4ac85..ab8f2be30c4 100644
--- a/lib/VM/JIT/arm64/JitHandlers.h
+++ b/lib/VM/JIT/arm64/JitHandlers.h
@@ -114,8 +114,7 @@ void *_jit_find_catch_target(
 [[noreturn]] void _sh_throw_register_stack_overflow(SHRuntime *shr);
 
 /// Call the closure stored in the outgoing registers of the current frame. The
-/// caller is responsible for setting up the callee closure, arg count, and
-/// new.target registers.
-SHLegacyValue _jit_dispatch_call(SHRuntime *, SHLegacyValue *frame);
+/// caller is responsible for setting up the outgoing registers.
+SHLegacyValue _jit_dispatch_call(SHRuntime *, SHLegacyValue *callTargetSHLV);
 
 } // namespace hermes::vm