Bug 1688033: Support FunApplyArgsObj in Warp r=jandem

As much as possible, I based `LApplyArgsObj` on `LApplyArrayGeneric`.

I wrote the tests here by looking at existing `apply` testcases in `jit-tests/tests/arguments` and writing similar test cases where the arguments object escapes.

Differential Revision: https://phabricator.services.mozilla.com/D104488

js/src/jit-test/tests/arguments/apply-args-obj-01.js

@@ -0,0 +1,18 @@
+function escape(x) { with ({}) {} }
+
+function foo() {
+    escape(arguments);
+    return bar.apply({}, arguments);
+}
+
+function bar(x,y) {
+    return x + y;
+}
+
+with ({}) {}
+
+var sum = 0;
+for (var i = 0; i < 100; i++) {
+    sum += foo(1,2);
+}
+assertEq(sum, 300);

js/src/jit-test/tests/arguments/apply-args-obj-02.js

@@ -0,0 +1,16 @@
+function foo() {
+    arguments[0] = 3;
+    return bar.apply({}, arguments);
+}
+
+function bar(x,y) {
+    return x + y;
+}
+
+with ({}) {}
+
+var sum = 0;
+for (var i = 0; i < 100; i++) {
+    sum += foo(1,2);
+}
+assertEq(sum, 500);

js/src/jit-test/tests/arguments/apply-args-obj-03.js

@@ -0,0 +1,20 @@
+function escape() { with ({}) {} }
+
+function foo(i) {
+    return i;
+}
+
+function bar(n) {
+    escape(arguments);
+    return foo.apply({}, arguments);
+}
+
+function baz(a, n) {
+    return bar(n);
+}
+
+var sum = 0;
+for (var i = 0; i < 10000; i++) {
+    sum += baz(0, 1);
+}
+assertEq(sum, 10000);

js/src/jit-test/tests/arguments/apply-args-obj-04.js

@@ -0,0 +1,21 @@
+var result;
+
+function g(a, b) {
+    with ({}) {}
+    result = a + b;
+}
+
+function escape() { with({}) {} }
+
+function f() {
+    escape(arguments);
+    for (var i = 0; i < 50; ++i) {
+	g.apply(this, arguments);
+    }
+}
+
+f(1, 2);
+assertEq(result, 3);
+
+f("");
+assertEq(result, "undefined");

js/src/jit/CodeGenerator.cpp

@@ -5690,6 +5690,35 @@ void CodeGenerator::emitPushArguments(LApplyArgsGeneric* apply,
   masm.pushValue(ToValue(apply, LApplyArgsGeneric::ThisIndex));
 }
 
+void CodeGenerator::emitPushArguments(LApplyArgsObj* apply,
+                                      Register extraStackSpace) {
+  // argc and argsObj are mapped to the same calltemp register.
+  MOZ_ASSERT(apply->getArgsObj() == apply->getArgc());
+
+  Register tmpArgc = ToRegister(apply->getTempObject());
+  Register argsObj = ToRegister(apply->getArgsObj());
+
+  // Load argc into tmpArgc.
+  Address lengthAddr(argsObj, ArgumentsObject::getInitialLengthSlotOffset());
+  masm.unboxInt32(lengthAddr, tmpArgc);
+  masm.rshift32(Imm32(ArgumentsObject::PACKED_BITS_COUNT), tmpArgc);
+
+  // Allocate space on the stack for arguments. This modifies extraStackSpace.
+  emitAllocateSpaceForApply(tmpArgc, extraStackSpace);
+
+  // Load arguments data
+  masm.loadPrivate(Address(argsObj, ArgumentsObject::getDataSlotOffset()),
+                   argsObj);
+  size_t argsSrcOffset = ArgumentsData::offsetOfArgs();
+
+  // This is the end of the lifetime of argsObj.
+  // After this call, the argsObj register holds the argument count instead.
+  emitPushArrayAsArguments(tmpArgc, argsObj, extraStackSpace, argsSrcOffset);
+
+  masm.addPtr(Imm32(sizeof(Value)), extraStackSpace);
+  masm.pushValue(ToValue(apply, LApplyArgsObj::ThisIndex));
+}
+
 void CodeGenerator::emitPushArrayAsArguments(Register tmpArgc,
                                              Register srcBaseAndArgc,
                                              Register scratch,
@@ -5810,16 +5839,17 @@ void CodeGenerator::emitApplyGeneric(T* apply) {
   Register objreg = ToRegister(apply->getTempObject());
   Register extraStackSpace = ToRegister(apply->getTempStackCounter());
 
-  // Holds the function nargs, computed in the invoker or (for ApplyArray and
-  // ConstructArray) in the argument pusher.
+  // Holds the function nargs, computed in the invoker or (for ApplyArray,
+  // ConstructArray, or ApplyArgsObj) in the argument pusher.
   Register argcreg = ToRegister(apply->getArgc());
 
   // Copy the arguments of the current function.
   //
-  // In the case of ApplyArray and ConstructArray, also compute argc: the argc
-  // register and the elements register are the same; argc must not be
-  // referenced before the call to emitPushArguments() and elements must not be
-  // referenced after it returns.
+  // In the case of ApplyArray, ConstructArray, or ApplyArgsObj, also
+  // compute argc. The argc register and the elements/argsObj register
+  // are the same; argc must not be referenced before the call to
+  // emitPushArguments() and elements/argsObj must not be referenced
+  // after it returns.
   //
   // In the case of ConstructArray, also overwrite newTarget with
   // extraStackSpace; newTarget must not be referenced after this point.
@@ -5986,6 +6016,18 @@ void CodeGenerator::visitApplyArgsGeneric(LApplyArgsGeneric* apply) {
   emitApplyGeneric(apply);
 }
 
+void CodeGenerator::visitApplyArgsObj(LApplyArgsObj* apply) {
+  Register argsObj = ToRegister(apply->getArgsObj());
+  Register temp = ToRegister(apply->getTempObject());
+
+  Label bail;
+  masm.loadArgumentsObjectLength(argsObj, temp, &bail);
+  masm.branch32(Assembler::Above, temp, Imm32(JIT_ARGS_LENGTH_MAX), &bail);
+  bailoutFrom(&bail, apply->snapshot());
+
+  emitApplyGeneric(apply);
+}
+
 void CodeGenerator::visitApplyArrayGeneric(LApplyArrayGeneric* apply) {
   LSnapshot* snapshot = apply->snapshot();
   Register tmp = ToRegister(apply->getTempObject());

js/src/jit/CodeGenerator.h

@@ -181,6 +181,7 @@ class CodeGenerator final : public CodeGeneratorSpecific {
   void emitPushArrayAsArguments(Register tmpArgc, Register srcBaseAndArgc,
                                 Register scratch, size_t argvSrcOffset);
   void emitPushArguments(LApplyArgsGeneric* apply, Register extraStackSpace);
+  void emitPushArguments(LApplyArgsObj* apply, Register extraStackSpace);
   void emitPushArguments(LApplyArrayGeneric* apply, Register extraStackSpace);
   void emitPushArguments(LConstructArrayGeneric* construct,
                          Register extraStackSpace);

js/src/jit/Lowering.cpp

@@ -461,6 +461,27 @@ void LIRGenerator::visitApplyArgs(MApplyArgs* apply) {
   assignSafepoint(lir, apply);
 }
 
+void LIRGenerator::visitApplyArgsObj(MApplyArgsObj* apply) {
+  MOZ_ASSERT(apply->getFunction()->type() == MIRType::Object);
+
+  // Assert if the return value is already erased.
+  static_assert(CallTempReg2 != JSReturnReg_Type);
+  static_assert(CallTempReg2 != JSReturnReg_Data);
+
+  LApplyArgsObj* lir = new (alloc()) LApplyArgsObj(
+      useFixedAtStart(apply->getFunction(), CallTempReg3),
+      useFixedAtStart(apply->getArgsObj(), CallTempReg0),
+      useBoxFixedAtStart(apply->getThis(), CallTempReg4, CallTempReg5),
+      tempFixed(CallTempReg1),   // object register
+      tempFixed(CallTempReg2));  // stack counter register
+
+  // Bailout is needed in the case of too many values in the arguments array.
+  assignSnapshot(lir, apply->bailoutKind());
+
+  defineReturn(lir, apply);
+  assignSafepoint(lir, apply);
+}
+
 void LIRGenerator::visitApplyArray(MApplyArray* apply) {
   MOZ_ASSERT(apply->getFunction()->type() == MIRType::Object);
 

js/src/jit/MIR.h

@@ -2533,6 +2533,40 @@ class MApplyArgs : public MTernaryInstruction,
   bool possiblyCalls() const override { return true; }
 };
 
+class MApplyArgsObj
+    : public MTernaryInstruction,
+      public MixPolicy<ObjectPolicy<0>, ObjectPolicy<1>, BoxPolicy<2>>::Data {
+  WrappedFunction* target_;
+  bool maybeCrossRealm_ = true;
+  bool ignoresReturnValue_ = false;
+
+  MApplyArgsObj(WrappedFunction* target, MDefinition* fun, MDefinition* argsObj,
+                MDefinition* thisArg)
+      : MTernaryInstruction(classOpcode, fun, argsObj, thisArg),
+        target_(target) {
+    MOZ_ASSERT(argsObj->type() == MIRType::Object);
+    setResultType(MIRType::Value);
+    setBailoutKind(BailoutKind::TooManyArguments);
+  }
+
+ public:
+  INSTRUCTION_HEADER(ApplyArgsObj)
+  TRIVIAL_NEW_WRAPPERS
+  NAMED_OPERANDS((0, getFunction), (1, getArgsObj), (2, getThis))
+
+  WrappedFunction* getSingleTarget() const { return target_; }
+
+  bool maybeCrossRealm() const { return maybeCrossRealm_; }
+  void setNotCrossRealm() { maybeCrossRealm_ = false; }
+
+  bool ignoresReturnValue() const { return ignoresReturnValue_; }
+  void setIgnoresReturnValue() { ignoresReturnValue_ = true; }
+
+  bool isConstructing() const { return false; }
+
+  bool possiblyCalls() const override { return true; }
+};
+
 // fun.apply(fn, array)
 class MApplyArray : public MTernaryInstruction,
                     public MixPolicy<ObjectPolicy<0>, BoxPolicy<2>>::Data {

js/src/jit/TypePolicy.cpp

@@ -1031,6 +1031,7 @@ bool ClampPolicy::adjustInputs(TempAllocator& alloc, MInstruction* ins) const {
   _(MixPolicy<ObjectPolicy<0>, BoxPolicy<2>>)                                 \
   _(MixPolicy<ObjectPolicy<0>, ObjectPolicy<1>, UnboxedInt32Policy<2>>)       \
   _(MixPolicy<ObjectPolicy<0>, ObjectPolicy<1>, ObjectPolicy<2>>)             \
+  _(MixPolicy<ObjectPolicy<0>, ObjectPolicy<1>, BoxPolicy<2>>)                \
   _(MixPolicy<StringPolicy<0>, UnboxedInt32Policy<1>, UnboxedInt32Policy<2>>) \
   _(MixPolicy<StringPolicy<0>, ObjectPolicy<1>, StringPolicy<2>>)             \
   _(MixPolicy<StringPolicy<0>, StringPolicy<1>, StringPolicy<2>>)             \

js/src/jit/WarpBuilderShared.h

@@ -51,6 +51,7 @@ class MOZ_STACK_CLASS CallInfo {
     Standard,
     Array,
     FunApplyMagicArgs,
+    FunApplyArgsObj,
   };
 
  private:

js/src/jit/WarpCacheIRTranspiler.cpp

@@ -243,6 +243,8 @@ class MOZ_RAII WarpCacheIRTranspiler : public WarpBuilderShared {
                                       CallFlags flags, CallKind kind);
   [[nodiscard]] bool emitFunApplyMagicArgs(WrappedFunction* wrappedTarget,
                                            CallFlags flags);
+  [[nodiscard]] bool emitFunApplyArgsObj(WrappedFunction* wrappedTarget,
+                                         CallFlags flags);
 
   MDefinition* convertWasmArg(MDefinition* arg, wasm::ValType::Kind kind);
 
@@ -3970,6 +3972,12 @@ bool WarpCacheIRTranspiler::updateCallInfo(MDefinition* callee,
         callInfo_->removeArg(0);
       }
       break;
+    case CallFlags::FunApplyArgsObj:
+      MOZ_ASSERT(!callInfo_->constructing());
+      MOZ_ASSERT(callInfo_->argFormat() == CallInfo::ArgFormat::Standard);
+
+      callInfo_->setArgFormat(CallInfo::ArgFormat::FunApplyArgsObj);
+      break;
     case CallFlags::FunApplyMagicArgs:
       MOZ_ASSERT(!callInfo_->constructing());
       MOZ_ASSERT(callInfo_->argFormat() == CallInfo::ArgFormat::Standard);
@@ -4114,6 +4122,9 @@ bool WarpCacheIRTranspiler::emitCallFunction(
     case CallInfo::ArgFormat::FunApplyMagicArgs: {
       return emitFunApplyMagicArgs(wrappedTarget, flags);
     }
+    case CallInfo::ArgFormat::FunApplyArgsObj: {
+      return emitFunApplyArgsObj(wrappedTarget, flags);
+    }
   }
   MOZ_CRASH("unreachable");
 }
@@ -4145,6 +4156,31 @@ bool WarpCacheIRTranspiler::emitFunApplyMagicArgs(
   return resumeAfter(apply);
 }
 
+bool WarpCacheIRTranspiler::emitFunApplyArgsObj(WrappedFunction* wrappedTarget,
+                                                CallFlags flags) {
+  MOZ_ASSERT(!callInfo_->constructing());
+  MOZ_ASSERT(!builder_->inlineCallInfo());
+
+  MDefinition* callee = callInfo_->thisArg();
+  MDefinition* thisArg = callInfo_->getArg(0);
+  MDefinition* argsObj = callInfo_->getArg(1);
+
+  MApplyArgsObj* apply =
+      MApplyArgsObj::New(alloc(), wrappedTarget, callee, argsObj, thisArg);
+
+  if (flags.isSameRealm()) {
+    apply->setNotCrossRealm();
+  }
+  if (callInfo_->ignoresReturnValue()) {
+    apply->setIgnoresReturnValue();
+  }
+
+  addEffectful(apply);
+  pushResult(apply);
+
+  return resumeAfter(apply);
+}
+
 #ifndef JS_SIMULATOR
 bool WarpCacheIRTranspiler::emitCallNativeFunction(ObjOperandId calleeId,
                                                    Int32OperandId argcId,

js/src/jit/shared/LIR-shared.h

@@ -1252,6 +1252,38 @@ class LApplyArgsGeneric
   const LDefinition* getTempStackCounter() { return getTemp(1); }
 };
 
+class LApplyArgsObj
+    : public LCallInstructionHelper<BOX_PIECES, BOX_PIECES + 2, 2> {
+ public:
+  LIR_HEADER(ApplyArgsObj)
+
+  LApplyArgsObj(const LAllocation& func, const LAllocation& argsObj,
+                const LBoxAllocation& thisv, const LDefinition& tmpObjReg,
+                const LDefinition& tmpCopy)
+      : LCallInstructionHelper(classOpcode) {
+    setOperand(0, func);
+    setOperand(1, argsObj);
+    setBoxOperand(ThisIndex, thisv);
+    setTemp(0, tmpObjReg);
+    setTemp(1, tmpCopy);
+  }
+
+  MApplyArgsObj* mir() const { return mir_->toApplyArgsObj(); }
+
+  bool hasSingleTarget() const { return getSingleTarget() != nullptr; }
+  WrappedFunction* getSingleTarget() const { return mir()->getSingleTarget(); }
+
+  const LAllocation* getFunction() { return getOperand(0); }
+  const LAllocation* getArgsObj() { return getOperand(1); }
+  // All registers are calltemps. argc is mapped to the same register as
+  // ArgsObj. argc becomes live as ArgsObj is dying.
+  const LAllocation* getArgc() { return getOperand(1); }
+  static const size_t ThisIndex = 2;
+
+  const LDefinition* getTempObject() { return getTemp(0); }
+  const LDefinition* getTempStackCounter() { return getTemp(1); }
+};
+
 class LApplyArrayGeneric
     : public LCallInstructionHelper<BOX_PIECES, BOX_PIECES + 2, 2> {
  public: