Bug 1355155 - Scalar replacement for arrow functions. r=nbp

js/src/jit/IonBuilder.cpp

@@ -12016,8 +12016,10 @@ IonBuilder::jsop_lambda_arrow(JSFunction* fun)
     MOZ_ASSERT(!fun->isNative());
 
     MDefinition* newTargetDef = current->pop();
+    MConstant* cst = MConstant::NewConstraintlessObject(alloc(), fun);
+    current->add(cst);
     MLambdaArrow* ins = MLambdaArrow::New(alloc(), constraints(), current->environmentChain(),
-                                          newTargetDef, fun);
+                                          newTargetDef, cst);
     current->add(ins);
     current->push(ins);
 

js/src/jit/MIR.h

@@ -8523,19 +8523,20 @@ class MLambda
 };
 
 class MLambdaArrow
-  : public MBinaryInstruction,
-    public MixPolicy<ObjectPolicy<0>, BoxPolicy<1>>::Data
+  : public MTernaryInstruction,
+    public Mix3Policy<ObjectPolicy<0>, BoxPolicy<1>, ObjectPolicy<2>>::Data
 {
     const LambdaFunctionInfo info_;
 
     MLambdaArrow(CompilerConstraintList* constraints, MDefinition* envChain,
-                 MDefinition* newTarget_, JSFunction* fun)
-      : MBinaryInstruction(envChain, newTarget_), info_(fun)
+                 MDefinition* newTarget, MConstant* cst)
+      : MTernaryInstruction(envChain, newTarget, cst),
+        info_(&cst->toObject().as<JSFunction>())
     {
         setResultType(MIRType::Object);
-        MOZ_ASSERT(!ObjectGroup::useSingletonForClone(fun));
-        if (!fun->isSingleton())
-            setResultTypeSet(MakeSingletonTypeSet(constraints, fun));
+        MOZ_ASSERT(!ObjectGroup::useSingletonForClone(info().fun));
+        if (!info().fun->isSingleton())
+            setResultTypeSet(MakeSingletonTypeSet(constraints, info().fun));
     }
 
   public:
@@ -8543,9 +8544,16 @@ class MLambdaArrow
     TRIVIAL_NEW_WRAPPERS
     NAMED_OPERANDS((0, environmentChain), (1, newTargetDef))
 
+    MConstant* functionOperand() const {
+        return getOperand(2)->toConstant();
+    }
     const LambdaFunctionInfo& info() const {
         return info_;
     }
+    MOZ_MUST_USE bool writeRecoverData(CompactBufferWriter& writer) const override;
+    bool canRecoverOnBailout() const override {
+        return true;
+    }
     bool appendRoots(MRootList& roots) const override {
         return info_.appendRoots(roots);
     }

js/src/jit/Recover.cpp

@@ -1477,6 +1477,35 @@ RLambda::recover(JSContext* cx, SnapshotIterator& iter) const
     return true;
 }
 
+bool
+MLambdaArrow::writeRecoverData(CompactBufferWriter& writer) const
+{
+    MOZ_ASSERT(canRecoverOnBailout());
+    writer.writeUnsigned(uint32_t(RInstruction::Recover_LambdaArrow));
+    return true;
+}
+
+RLambdaArrow::RLambdaArrow(CompactBufferReader& reader)
+{
+}
+
+bool
+RLambdaArrow::recover(JSContext* cx, SnapshotIterator& iter) const
+{
+    RootedObject scopeChain(cx, &iter.read().toObject());
+    RootedValue newTarget(cx, iter.read());
+    RootedFunction fun(cx, &iter.read().toObject().as<JSFunction>());
+
+    JSObject* resultObject = js::LambdaArrow(cx, fun, scopeChain, newTarget);
+    if (!resultObject)
+        return false;
+
+    RootedValue result(cx);
+    result.setObject(*resultObject);
+    iter.storeInstructionResult(result);
+    return true;
+}
+
 bool
 MSimdBox::writeRecoverData(CompactBufferWriter& writer) const
 {

js/src/jit/Recover.h

@@ -106,6 +106,7 @@ namespace jit {
     _(NewDerivedTypedObject)                    \
     _(CreateThisWithTemplate)                   \
     _(Lambda)                                   \
+    _(LambdaArrow)                              \
     _(SimdBox)                                  \
     _(ObjectState)                              \
     _(ArrayState)                               \
@@ -626,6 +627,14 @@ class RLambda final : public RInstruction
     MOZ_MUST_USE bool recover(JSContext* cx, SnapshotIterator& iter) const override;
 };
 
+class RLambdaArrow final : public RInstruction
+{
+  public:
+    RINSTRUCTION_HEADER_NUM_OP_(LambdaArrow, 3)
+
+    MOZ_MUST_USE bool recover(JSContext* cx, SnapshotIterator& iter) const override;
+};
+
 class RSimdBox final : public RInstruction
 {
   private:

js/src/jit/ScalarReplacement.cpp

@@ -101,8 +101,9 @@ IsObjectEscaped(MInstruction* ins, JSObject* objDefault = nullptr);
 // Returns False if the lambda is not escaped and if it is optimizable by
 // ScalarReplacementOfObject.
 static bool
-IsLambdaEscaped(MLambda* lambda, JSObject* obj)
+IsLambdaEscaped(MInstruction* lambda, JSObject* obj)
 {
+    MOZ_ASSERT(lambda->isLambda() || lambda->isLambdaArrow());
     JitSpewDef(JitSpew_Escape, "Check lambda\n", lambda);
     JitSpewIndent spewIndent(JitSpew_Escape);
 
@@ -242,10 +243,10 @@ IsObjectEscaped(MInstruction* ins, JSObject* objDefault)
             break;
           }
 
-          case MDefinition::Op_Lambda: {
-            MLambda* lambda = def->toLambda();
-            if (IsLambdaEscaped(lambda, obj)) {
-                JitSpewDef(JitSpew_Escape, "is indirectly escaped by\n", lambda);
+          case MDefinition::Op_Lambda:
+          case MDefinition::Op_LambdaArrow: {
+            if (IsLambdaEscaped(def->toInstruction(), obj)) {
+                JitSpewDef(JitSpew_Escape, "is indirectly escaped by\n", def);
                 return true;
             }
             break;
@@ -312,6 +313,7 @@ class ObjectMemoryView : public MDefinitionVisitorDefaultNoop
     void visitGuardShape(MGuardShape* ins);
     void visitFunctionEnvironment(MFunctionEnvironment* ins);
     void visitLambda(MLambda* ins);
+    void visitLambdaArrow(MLambdaArrow* ins);
     void visitStoreUnboxedScalar(MStoreUnboxedScalar* ins);
     void visitLoadUnboxedScalar(MLoadUnboxedScalar* ins);
     void visitStoreUnboxedObjectOrNull(MStoreUnboxedObjectOrNull* ins);
@@ -480,7 +482,7 @@ ObjectMemoryView::assertSuccess()
 
         // The only remaining uses would be removed by DCE, which will also
         // recover the object on bailouts.
-        MOZ_ASSERT(def->isSlots() || def->isLambda());
+        MOZ_ASSERT(def->isSlots() || def->isLambda() || def->isLambdaArrow());
         MOZ_ASSERT(!def->hasDefUses());
     }
 }
@@ -642,8 +644,15 @@ ObjectMemoryView::visitFunctionEnvironment(MFunctionEnvironment* ins)
 {
     // Skip function environment which are not aliases of the NewCallObject.
     MDefinition* input = ins->input();
-    if (!input->isLambda() || input->toLambda()->environmentChain() != obj_)
+    if (input->isLambda()) {
+        if (input->toLambda()->environmentChain() != obj_)
+            return;
+    } else if (input->isLambdaArrow()) {
+        if (input->toLambdaArrow()->environmentChain() != obj_)
+            return;
+    } else {
         return;
+    }
 
     // Replace the function environment by the scope chain of the lambda.
     ins->replaceAllUsesWith(obj_);
@@ -663,6 +672,15 @@ ObjectMemoryView::visitLambda(MLambda* ins)
     ins->setIncompleteObject();
 }
 
+void
+ObjectMemoryView::visitLambdaArrow(MLambdaArrow* ins)
+{
+    if (ins->environmentChain() != obj_)
+        return;
+
+    ins->setIncompleteObject();
+}
+
 static size_t
 GetOffsetOf(MDefinition* index, size_t width, int32_t baseOffset)
 {

js/src/jit/shared/CodeGenerator-shared.cpp

@@ -412,8 +412,9 @@ CodeGeneratorShared::encodeAllocation(LSnapshot* snapshot, MDefinition* mir,
 
         // Lambda should have a default value readable for iterating over the
         // inner frames.
-        if (mir->isLambda()) {
-            MConstant* constant = mir->toLambda()->functionOperand();
+        if (mir->isLambda() || mir->isLambdaArrow()) {
+            MConstant* constant = mir->isLambda() ? mir->toLambda()->functionOperand()
+                                                  : mir->toLambdaArrow()->functionOperand();
             uint32_t cstIndex;
             masm.propagateOOM(graph.addConstantToPool(constant->toJSValue(), &cstIndex));
             alloc = RValueAllocation::RecoverInstruction(index, cstIndex);