Save and restore generator locals in bulk with dedicated bytecodes

JSTests/stress/generator-save-restore-locals.js

@@ -0,0 +1,175 @@
+//@ runDefault
+//@ runDefault("--useGeneratorBulkSaveRestore=0")
+//@ runDefault("--useJIT=0")
+//@ runDefault("--useDFGJIT=0")
+//@ runDefault("--useConcurrentJIT=0", "--thresholdForJITAfterWarmUp=10", "--thresholdForOptimizeAfterWarmUp=20", "--thresholdForFTLOptimizeAfterWarmUp=50")
+//@ runBytecodeCache
+
+function shouldBe(actual, expected) {
+    if (actual !== expected)
+        throw new Error(`bad value: ${actual}, expected ${expected}`);
+}
+
+// Mixed-type locals, resumed with next/throw.
+function* mixed() {
+    let a = 1, b = 2.5, c = "s", d = { v: 3 }, e;
+    try {
+        a += yield a;
+        b += yield b;
+        e = yield c + d.v;
+    } catch (error) {
+        return [a, b, c, d.v, e, error];
+    }
+    return [a, b, c, d.v, e];
+}
+{
+    let iterator = mixed();
+    shouldBe(iterator.next().value, 1);
+    shouldBe(iterator.next(10).value, 2.5);
+    shouldBe(iterator.next(0.5).value, "s3");
+    shouldBe(JSON.stringify(iterator.throw("boom").value), '[11,3,"s",3,null,"boom"]');
+}
+
+// Live sets of various sizes: 70 needs an out-of-line bit vector, 300 needs wide operands.
+function makeWide(count) {
+    let declarations = [];
+    let names = [];
+    for (let i = 0; i < count; ++i) {
+        declarations.push(`let v${i} = ${i};`);
+        names.push(`v${i}`);
+    }
+    return new Function(`
+        return function* wide() {
+            ${declarations.join("\n")}
+            yield 0;
+            ${names.map((name) => `${name} += 1;`).join("\n")}
+            yield 1;
+            return ${names.join(" + ")};
+        }`)();
+}
+for (let count of [3, 9, 70, 300]) {
+    let wide = makeWide(count);
+    let expected = 0;
+    for (let i = 0; i < count; ++i)
+        expected += i + 1;
+    for (let round = 0; round < 3; ++round) {
+        let iterator = wide();
+        shouldBe(iterator.next().value, 0);
+        shouldBe(iterator.next().value, 1);
+        shouldBe(iterator.next().value, expected);
+    }
+}
+
+// Locals under TDZ hold the empty value across a yield.
+function* tdz(flag) {
+    yield 1;
+    if (flag) {
+        let later = 5;
+        yield later;
+    }
+    let after = 7;
+    yield after;
+    {
+        yield 2;
+        let x = 9;
+        yield x;
+    }
+}
+{
+    let iterator = tdz(false);
+    shouldBe(iterator.next().value, 1);
+    shouldBe(iterator.next().value, 7);
+    shouldBe(iterator.next().value, 2);
+    shouldBe(iterator.next().value, 9);
+}
+
+// Saved locals share the lexical environment with captured variables.
+function* captured() {
+    let counter = 0;
+    let local = 100;
+    const bump = () => ++counter;
+    yield bump();
+    local += counter;
+    yield bump() + local;
+    return [counter, local];
+}
+{
+    let iterator = captured();
+    shouldBe(iterator.next().value, 1);
+    shouldBe(iterator.next().value, 103);
+    shouldBe(JSON.stringify(iterator.next().value), "[2,101]");
+}
+
+// Nothing live across the yields.
+function* empty() {
+    yield 1;
+    yield 2;
+}
+{
+    let iterator = empty();
+    shouldBe(iterator.next().value, 1);
+    shouldBe(iterator.next().value, 2);
+    shouldBe(iterator.next().done, true);
+}
+
+// Async function, async generator, and yield*.
+async function asyncFunction(promise) {
+    let a = 1, b = [1, 2, 3], c = "x";
+    a += await promise;
+    for (const v of b)
+        c += await v;
+    return a + c;
+}
+async function* asyncGenerator() {
+    let accumulator = 0;
+    for (let i = 0; i < 3; ++i) {
+        accumulator += await i;
+        yield accumulator;
+    }
+    yield* empty();
+    return accumulator;
+}
+function* delegating() {
+    return yield* mixed();
+}
+let asyncDone = false;
+(async () => {
+    shouldBe(await asyncFunction(Promise.resolve(10)), "11x123");
+    let values = [];
+    for await (const value of asyncGenerator())
+        values.push(value);
+    shouldBe(JSON.stringify(values), "[0,1,3,1,2]");
+    asyncDone = true;
+})();
+{
+    let iterator = delegating();
+    iterator.next();
+    iterator.next(1);
+    iterator.next(1);
+    shouldBe(JSON.stringify(iterator.next().value), '[2,3.5,"s",3,null]');
+}
+
+// Tier up with int32 locals, then resume with double and string values.
+function* counting(count, start) {
+    let i = 0, accumulator = start, object = { f: 1 };
+    while (i < count) {
+        accumulator = accumulator + object.f;
+        yield accumulator;
+        ++i;
+    }
+    return accumulator;
+}
+function drive(start, count) {
+    let iterator = counting(count, start);
+    let last;
+    for (let result = iterator.next(); !result.done; result = iterator.next())
+        last = result.value;
+    return last;
+}
+for (let i = 0; i < testLoopCount; ++i)
+    shouldBe(drive(0, 10), 10);
+shouldBe(drive(0.5, 10), 10.5);
+shouldBe(drive("s", 3), "s111");
+
+drainMicrotasks();
+shouldBe(asyncDone, true);

Source/JavaScriptCore/bytecode/BytecodeGeneratorification.cpp

@@ -38,6 +38,7 @@
 #include "StrongInlines.h"
 #include "UnlinkedCodeBlockGenerator.h"
 #include "UnlinkedMetadataTableInlines.h"
+#include <wtf/HashMap.h>
 
 namespace JSC {
 
@@ -166,6 +167,9 @@ class BytecodeGeneratorification {
         return storage;
     }
 
+    void emitBulkSaveAndRestore(BytecodeRewriter&);
+    void emitPerLocalSaveAndRestore(BytecodeRewriter&);
+
     BytecodeGenerator& m_bytecodeGenerator;
     JSInstructionStream::Offset m_enterPoint;
     std::optional<GeneratorFrameData> m_generatorFrameData;
@@ -204,7 +208,6 @@ void BytecodeGeneratorification::run()
 {
     // We calculate the liveness at each merge point. This gives us the information which registers should be saved and resumed conservatively.
 
-    VM& vm = m_bytecodeGenerator.vm();
     {
         GeneratorLivenessAnalysis pass(*this);
         pass.run(m_codeBlock, m_instructions);
@@ -231,6 +234,89 @@ void BytecodeGeneratorification::run()
         });
     }
 
+    if (Options::useGeneratorBulkSaveRestore())
+        emitBulkSaveAndRestore(rewriter);
+    else
+        emitPerLocalSaveAndRestore(rewriter);
+
+    if (m_generatorFrameData) {
+        auto instruction = m_instructions.at(m_generatorFrameData->m_point);
+        rewriter.replaceBytecodeWithFragment(instruction, [&] (BytecodeRewriter::Fragment& fragment) {
+            if (!m_generatorFrameSymbolTable->scopeSize()) {
+                // This will cause us to put jsUndefined() into the generator frame's scope value.
+                fragment.appendInstruction<OpMov>(m_generatorFrameData->m_dst, m_generatorFrameData->m_initialValue);
+            } else
+                fragment.appendInstruction<OpCreateLexicalEnvironment>(m_generatorFrameData->m_dst, m_generatorFrameData->m_scope, m_generatorFrameData->m_symbolTable, m_generatorFrameData->m_initialValue);
+        });
+    }
+
+    rewriter.execute();
+}
+
+void BytecodeGeneratorification::emitBulkSaveAndRestore(BytecodeRewriter& rewriter)
+{
+    BitVector savedLocals;
+    for (const YieldData& data : m_yields) {
+        data.liveness.forEachSetBit([&](size_t index) {
+            savedLocals.set(index);
+        });
+    }
+
+    unsigned firstScopeOffset = m_generatorFrameSymbolTable->scopeSize();
+    for (unsigned i = 0, count = savedLocals.bitCount(); i < count; ++i) {
+        ScopeOffset scopeOffset = m_generatorFrameSymbolTable->takeNextScopeOffset(NoLockingNecessary);
+        ASSERT_UNUSED(scopeOffset, scopeOffset.offset() == firstScopeOffset + i);
+    }
+
+    unsigned numberOfLocalBits = savedLocals.size();
+    UncheckedKeyHashMap<BitVector, unsigned> bitVectorIndices;
+    auto addBitVectorConstant = [&](const BitVector& bitVector) {
+        return bitVectorIndices.ensure(bitVector, [&] {
+            return m_codeBlock->addBitVector(BitVector(bitVector));
+        }).iterator->value;
+    };
+    unsigned savedLocalsIndex = savedLocals.isEmpty() ? 0 : addBitVectorConstant(savedLocals);
+
+    for (const YieldData& data : m_yields) {
+        VirtualRegister scope = virtualRegisterForArgumentIncludingThis(static_cast<int32_t>(JSGenerator::Argument::Frame));
+        auto instruction = m_instructions.at(data.point);
+
+        if (data.liveness.isEmpty()) {
+            rewriter.insertFragmentBefore(instruction, [&] (BytecodeRewriter::Fragment& fragment) {
+                fragment.appendInstruction<OpRet>(data.argument);
+            });
+            rewriter.replaceBytecodeWithFragment(instruction, [&] (BytecodeRewriter::Fragment&) { });
+            continue;
+        }
+
+        BitVector liveLocals;
+        liveLocals.ensureSize(numberOfLocalBits);
+        data.liveness.forEachSetBit([&](size_t index) {
+            liveLocals.quickSet(index);
+        });
+        unsigned liveLocalsIndex = addBitVectorConstant(liveLocals);
+
+        unsigned firstValueProfile = m_bytecodeGenerator.nextValueProfileIndex();
+        for (unsigned i = 1, count = data.liveness.bitCount(); i < count; ++i)
+            m_bytecodeGenerator.nextValueProfileIndex();
+
+        // Emit save sequence.
+        rewriter.insertFragmentBefore(instruction, [&] (BytecodeRewriter::Fragment& fragment) {
+            fragment.appendInstruction<OpSaveGeneratorLocals>(scope, liveLocalsIndex, savedLocalsIndex, firstScopeOffset);
+            // Insert op_ret just after save sequence.
+            fragment.appendInstruction<OpRet>(data.argument);
+        });
+
+        // Emit resume sequence.
+        rewriter.replaceBytecodeWithFragment(instruction, [&] (BytecodeRewriter::Fragment& fragment) {
+            fragment.appendInstruction<OpRestoreGeneratorLocals>(scope, liveLocalsIndex, savedLocalsIndex, firstScopeOffset, firstValueProfile);
+        });
+    }
+}
+
+void BytecodeGeneratorification::emitPerLocalSaveAndRestore(BytecodeRewriter& rewriter)
+{
+    VM& vm = m_bytecodeGenerator.vm();
     for (const YieldData& data : m_yields) {
         VirtualRegister scope = virtualRegisterForArgumentIncludingThis(static_cast<int32_t>(JSGenerator::Argument::Frame));
 
@@ -273,19 +359,6 @@ void BytecodeGeneratorification::run()
             });
         });
     }
-
-    if (m_generatorFrameData) {
-        auto instruction = m_instructions.at(m_generatorFrameData->m_point);
-        rewriter.replaceBytecodeWithFragment(instruction, [&] (BytecodeRewriter::Fragment& fragment) {
-            if (!m_generatorFrameSymbolTable->scopeSize()) {
-                // This will cause us to put jsUndefined() into the generator frame's scope value.
-                fragment.appendInstruction<OpMov>(m_generatorFrameData->m_dst, m_generatorFrameData->m_initialValue);
-            } else
-                fragment.appendInstruction<OpCreateLexicalEnvironment>(m_generatorFrameData->m_dst, m_generatorFrameData->m_scope, m_generatorFrameData->m_symbolTable, m_generatorFrameData->m_initialValue);
-        });
-    }
-
-    rewriter.execute();
 }
 
 void performGeneratorification(BytecodeGenerator& bytecodeGenerator, UnlinkedCodeBlockGenerator* codeBlock, JSInstructionStreamWriter& instructions, SymbolTable* generatorFrameSymbolTable, int generatorFrameSymbolTableIndex)

Source/JavaScriptCore/bytecode/BytecodeGeneratorification.h

@@ -26,12 +26,13 @@
 
 #pragma once
 
+#include <wtf/BitVector.h>
+
 namespace JSC {
 
 class BytecodeGenerator;
 class SymbolTable;
 class UnlinkedCodeBlockGenerator;
-class SymbolTable;
 
 struct JSOpcodeTraits;
 template<typename> struct BaseInstruction;
@@ -41,4 +42,15 @@ using JSInstructionStreamWriter = InstructionStreamWriter<JSInstruction>;
 
 void performGeneratorification(BytecodeGenerator&, UnlinkedCodeBlockGenerator*, JSInstructionStreamWriter&, SymbolTable* generatorFrameSymbolTable, int generatorFrameSymbolTableIndex);
 
+template<typename Functor>
+void forEachLiveGeneratorLocal(const BitVector& liveLocals, const BitVector& savedLocals, const Functor& functor)
+{
+    unsigned slot = 0;
+    savedLocals.forEachSetBit([&](size_t index) {
+        if (liveLocals.quickGet(index))
+            functor(index, slot);
+        ++slot;
+    });
+}
+
 } // namespace JSC

Source/JavaScriptCore/bytecode/BytecodeList.rb

@@ -1176,6 +1176,23 @@
         argument: VirtualRegister,
     }
 
+op :save_generator_locals,
+    args: {
+        scope: VirtualRegister,
+        liveLocals: unsigned,
+        savedLocals: unsigned,
+        firstScopeOffset: unsigned,
+    }
+
+op :restore_generator_locals,
+    args: {
+        scope: VirtualRegister,
+        liveLocals: unsigned,
+        savedLocals: unsigned,
+        firstScopeOffset: unsigned,
+        valueProfile: unsigned,
+    }
+
 op :check_traps
 
 op :log_shadow_chicken_prologue,

Source/JavaScriptCore/bytecode/BytecodeUseDef.cpp

@@ -292,6 +292,8 @@ void computeUsesForBytecodeIndexImpl(const JSInstruction* instruction, Checkpoin
     USES(OpPutInternalField, base, value)
 
     USES(OpYield, argument)
+    USES(OpSaveGeneratorLocals, scope)
+    USES(OpRestoreGeneratorLocals, scope)
 
     USES(OpEnumeratorNext, mode, index, base, enumerator)
     USES(OpEnumeratorGetByVal, base, mode, propertyName, index, enumerator)
@@ -459,6 +461,8 @@ void computeDefsForBytecodeIndexImpl(unsigned numVars, const JSInstruction* inst
     case op_log_shadow_chicken_prologue:
     case op_log_shadow_chicken_tail:
     case op_yield:
+    case op_save_generator_locals:
+    case op_restore_generator_locals:
     case op_nop:
     case op_unreachable:
     case op_super_sampler_begin: