From fd99ca52128b1b48b2c588f3eed036812d1af703 Mon Sep 17 00:00:00 2001
From: Tzvetan Mikov <tmikov@meta.com>
Date: Wed, 28 Aug 2024 08:14:28 -0700
Subject: [PATCH] jit: implement constants except strings and bigint

Summary: .

Reviewed By: avp

Differential Revision: D61764349

fbshipit-source-id: 40e6c1e3abcfcd4bb18e8fa8e82576c4673340a5
---
 lib/VM/JIT/arm64/JIT.cpp        | 27 ++++++++++++++++++++++++---
 lib/VM/JIT/arm64/JitEmitter.cpp | 23 +++++++++++++++++++++--
 lib/VM/JIT/arm64/JitEmitter.h   |  3 ++-
 test/jit/consts.js              | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 6 deletions(-)
 create mode 100644 test/jit/consts.js

diff --git a/lib/VM/JIT/arm64/JIT.cpp b/lib/VM/JIT/arm64/JIT.cpp
index a0c749fd73c..60b6037b0e1 100644
--- a/lib/VM/JIT/arm64/JIT.cpp
+++ b/lib/VM/JIT/arm64/JIT.cpp
@@ -102,14 +102,35 @@ JITCompiledFunctionPtr JITContext::compileImpl(
           ip = NEXTINST(LoadParam);
           break;
         case inst::OpCode::LoadConstZero:
-          em.loadConstUInt8(FR(ip->iLoadConstZero.op1), 0);
+          em.loadConstDouble(FR(ip->iLoadConstZero.op1), 0, "Zero");
           ip = NEXTINST(LoadConstZero);
           break;
         case inst::OpCode::LoadConstUInt8:
-          em.loadConstUInt8(
-              FR(ip->iLoadConstUInt8.op1), ip->iLoadConstUInt8.op2);
+          em.loadConstDouble(
+              FR(ip->iLoadConstUInt8.op1), ip->iLoadConstUInt8.op2, "UInt8");
           ip = NEXTINST(LoadConstUInt8);
           break;
+        case inst::OpCode::LoadConstInt:
+          em.loadConstDouble(
+              FR(ip->iLoadConstInt.op1), ip->iLoadConstInt.op2, "Int");
+          ip = NEXTINST(LoadConstInt);
+          break;
+        case inst::OpCode::LoadConstDouble:
+          em.loadConstDouble(
+              FR(ip->iLoadConstDouble.op1), ip->iLoadConstDouble.op2, "Double");
+          ip = NEXTINST(LoadConstDouble);
+          break;
+#define LOAD_CONST(NAME, val, type)                                     \
+  case inst::OpCode::LoadConst##NAME:                                   \
+    em.loadConstBits64(FR(ip->iLoadConst##NAME.op1), val, type, #NAME); \
+    ip = NEXTINST(LoadConst##NAME);                                     \
+    break;
+          LOAD_CONST(Empty, _sh_ljs_empty().raw, FRType::Unknown);
+          LOAD_CONST(Undefined, _sh_ljs_undefined().raw, FRType::Unknown);
+          LOAD_CONST(Null, _sh_ljs_null().raw, FRType::Unknown);
+          LOAD_CONST(True, _sh_ljs_bool(true).raw, FRType::Bool);
+          LOAD_CONST(False, _sh_ljs_bool(false).raw, FRType::Bool);
+#undef LOAD_CONST
 
         case inst::OpCode::Mov:
           em.mov(FR(ip->iMov.op1), FR(ip->iMov.op2));
diff --git a/lib/VM/JIT/arm64/JitEmitter.cpp b/lib/VM/JIT/arm64/JitEmitter.cpp
index 5b1c0e442a6..10d79aaae94 100644
--- a/lib/VM/JIT/arm64/JitEmitter.cpp
+++ b/lib/VM/JIT/arm64/JitEmitter.cpp
@@ -720,8 +720,8 @@ void Emitter::loadParam(FR frRes, uint32_t paramIndex) {
   frUpdatedWithHWReg(frRes, hwRes);
 }
 
-void Emitter::loadConstUInt8(FR frRes, uint8_t val) {
-  comment("// LoadConstUInt8 r%u, %u", frRes.index(), val);
+void Emitter::loadConstDouble(FR frRes, double val, const char *name) {
+  comment("// LoadConst%s r%u, %f", name, frRes.index(), val);
   HWReg hwRes{};
 
   if (val == 0) {
@@ -745,6 +745,25 @@ void Emitter::loadConstUInt8(FR frRes, uint8_t val) {
   }
   frUpdatedWithHWReg(frRes, hwRes, FRType::Number);
 }
+void Emitter::loadConstBits64(
+    FR frRes,
+    uint64_t bits,
+    FRType type,
+    const char *name) {
+  comment(
+      "// LoadConst%s r%u, %llu",
+      name,
+      frRes.index(),
+      (unsigned long long)bits);
+  HWReg hwRes = getOrAllocFRInGpX(frRes, false);
+
+  if (isCheapConst(bits)) {
+    a.mov(hwRes.a64GpX(), bits);
+  } else {
+    a.ldr(hwRes.a64GpX(), a64::Mem(roDataLabel_, uint64Const(bits, name)));
+  }
+  frUpdatedWithHWReg(frRes, hwRes, type);
+}
 
 void Emitter::toNumber(FR frRes, FR frInput) {
   comment("// %s r%u, r%u", "toNumber", frRes.index(), frInput.index());
diff --git a/lib/VM/JIT/arm64/JitEmitter.h b/lib/VM/JIT/arm64/JitEmitter.h
index d474bdb170a..702cffd0871 100644
--- a/lib/VM/JIT/arm64/JitEmitter.h
+++ b/lib/VM/JIT/arm64/JitEmitter.h
@@ -368,7 +368,8 @@ class Emitter {
   void ret(FR frValue);
   void mov(FR frRes, FR frInput, bool logComment = true);
   void loadParam(FR frRes, uint32_t paramIndex);
-  void loadConstUInt8(FR frRes, uint8_t val);
+  void loadConstDouble(FR frRes, double val, const char *name);
+  void loadConstBits64(FR frRes, uint64_t val, FRType type, const char *name);
   void toNumber(FR frRes, FR frInput);
 
 #define DECL_BINOP(methodName, forceNum, commentStr, slowCall, a64body) \
diff --git a/test/jit/consts.js b/test/jit/consts.js
new file mode 100644
index 00000000000..1acc94531f1
--- /dev/null
+++ b/test/jit/consts.js
@@ -0,0 +1,33 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// RUN: %hermes -Xforce-jit -Xdump-jitcode %s | %FileCheck --match-full-lines %s
+// REQUIRES: jit
+
+function constants(o) {
+    o.zero = 0;
+    o.uint8 = 5;
+    o.int = 100_000_000;
+    o.d = 3.14;
+    o.undefined = void 0;
+    o.null = null;
+    o.true = true;
+    o.false = false;
+    return o;
+}
+
+print(JSON.stringify(constants({}), undefined, 2));
+// CHECK: JIT successfully compiled FunctionID 1, 'constants'
+// CHECK-NEXT: {
+// CHECK-NEXT:   "zero": 0,
+// CHECK-NEXT:   "uint8": 5,
+// CHECK-NEXT:   "int": 100000000,
+// CHECK-NEXT:   "d": 3.14,
+// CHECK-NEXT:   "null": null,
+// CHECK-NEXT:   "true": true,
+// CHECK-NEXT:   "false": false
+// CHECK-NEXT: }