continue a in a: b: for goes on to the next iteration; labelled for-of no longer crashes

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -1650,7 +1650,12 @@ namespace mlirgen
         // the labels around a function are not the body's to jump to
         auto outerLabels = std::move(activeLabels);
         activeLabels.clear();
-        auto restoreLabels = llvm::make_scope_exit([&]() { activeLabels = std::move(outerLabels); });
+        auto outerLoopLabelAliases = std::move(loopLabelAliases);
+        loopLabelAliases.clear();
+        auto restoreLabels = llvm::make_scope_exit([&]() {
+            activeLabels = std::move(outerLabels);
+            loopLabelAliases = std::move(outerLoopLabelAliases);
+        });
 
         auto location = loc(functionLikeDeclarationBaseAST);
 

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -753,6 +753,11 @@ class MLIRGenImpl
         // outwards instead, but the loop sites clear `label` before storing it, so a labelled
         // loop's context holds an empty one too - and `continue cont1` then relies on the outer
         // loop matching that empty label. 02disposable.ts is the case that proves it.
+        // (The `for` site now takes the label before its initializer and stores "" outright.)
+        // Known gap: a labelled `continue` written inside an `if` in an inner loop stops at the
+        // inner loop's body, so the outer body's `using` is not disposed. Carrying real labels
+        // needs the loop context, which stands in for the enclosing body scope in this walk, to
+        // keep that scope's label apart from the label of the loop it starts.
         auto isTargetLoop = genContext->isLoopBodyScope && genContext->loopLabel == loopLabel;
         return !isTargetLoop;
     }
@@ -14434,6 +14439,16 @@ class MLIRGenImpl
 
 private:
     std::string label;
+    // `a: b: for (...)`: the outer labels of a chain that ends in a loop, each mapped to the innermost
+    // one, the label the loop carries (see mlirGen(LabeledStatement))
+    llvm::StringMap<std::string> loopLabelAliases;
+
+    std::string resolveLoopLabelAlias(std::string label)
+    {
+        auto it = loopLabelAliases.find(label);
+        return it != loopLabelAliases.end() ? it->second : label;
+    }
+
     mlir::Block* tempEntryBlock;
     mlir::ModuleOp tempModule;
     mlir_ts::FuncOp tempFuncOp;

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -798,6 +798,10 @@ namespace mlirgen
 
         auto location = loc(forStatementAST);
 
+        // the loop's label, taken before the initializer: generating it can generate a loop of its own
+        auto loopLabel = label;
+        label = "";
+
         // initializer
         // TODO: why do we have ForInitialier
         if (isExpression(forStatementAST->initializer))
@@ -824,15 +828,15 @@ namespace mlirgen
         SmallVector<mlir::Value, 0> operands;
 
         auto forOp = builder.create<mlir_ts::ForOp>(location, types, operands);
-        if (!label.empty())
+        if (!loopLabel.empty())
         {
-            forOp->setAttr(LABEL_ATTR_NAME, builder.getStringAttr(label));
-            label = "";
+            forOp->setAttr(LABEL_ATTR_NAME, builder.getStringAttr(loopLabel));
         }
 
         GenContext loopGenContext(genContext);
         loopGenContext.isLoop = true;
-        loopGenContext.loopLabel = label;
+        // empty, as at the other loop sites: see scopeExitContinuesOutwards
+        loopGenContext.loopLabel = "";
 
         /*auto *cond =*/builder.createBlock(&forOp.getCond(), {}, types);
         /*auto *body =*/builder.createBlock(&forOp.getBody(), {}, types);
@@ -950,6 +954,11 @@ namespace mlirgen
     {
         auto location = loc(forOfStatementAST);
 
+        // the loop's label, handed back to the `for` the loop is made into: generating the iterated value or its
+        // iterator can generate functions whose own loops would take it
+        auto loopLabel = label;
+        label = "";
+
         auto result = mlirGen(forOfStatementAST->expression, genContext);
         EXIT_IF_FAILED_OR_NO_VALUE(result)
         auto exprValue = V(result);
@@ -971,13 +980,15 @@ namespace mlirgen
             auto propertyType = evaluateProperty(location, exprValue, ITERATOR_NEXT, genContext);
             if (propertyType)
             {
+                label = loopLabel;
                 if (mlir::succeeded(mlirGenES2015(forOfStatementAST, exprValue, genContext)))
                 {
                     return mlir::success();
                 }
             }
         }
 
+        label = loopLabel;
         return mlirGenES3(forOfStatementAST, exprValue, genContext);
     }
 
@@ -994,18 +1005,42 @@ namespace mlirgen
             return mlir::success();
         }
 
-        auto noLabelOp = kind == SyntaxKind::WhileStatement || kind == SyntaxKind::DoStatement ||
-                         kind == SyntaxKind::ForStatement || kind == SyntaxKind::ForInStatement ||
-                         kind == SyntaxKind::ForOfStatement;
+        auto isLoop = [](SyntaxKind kind) {
+            return kind == SyntaxKind::WhileStatement || kind == SyntaxKind::DoStatement ||
+                   kind == SyntaxKind::ForStatement || kind == SyntaxKind::ForInStatement ||
+                   kind == SyntaxKind::ForOfStatement;
+        };
 
         activeLabels.push_back(label);
         auto popLabel = llvm::make_scope_exit([&]() { activeLabels.pop_back(); });
 
-        if (noLabelOp)
+        if (isLoop(kind))
         {
             return mlirGen(labeledStatementAST->statement, genContext);
         }
 
+        // `a: b: for (...)`: both labels name the loop, which carries only the innermost one. A LabelOp
+        // for `a` would make `continue a` restart its region - the loop from its initializer, skipping
+        // the incrementor - so `a` becomes another name for `b` instead.
+        if (kind == SyntaxKind::LabeledStatement)
+        {
+            auto innermost = labeledStatementAST->statement.as<LabeledStatement>();
+            while ((SyntaxKind)innermost->statement == SyntaxKind::LabeledStatement)
+            {
+                innermost = innermost->statement.as<LabeledStatement>();
+            }
+
+            if (isLoop((SyntaxKind)innermost->statement))
+            {
+                // a function generated inside the loop starts with none (see mlirGenFunctionBody's activeLabels)
+                auto alias = label;
+                loopLabelAliases[alias] = MLIRHelper::getName(innermost->label);
+                auto res = mlirGen(labeledStatementAST->statement, genContext);
+                loopLabelAliases.erase(alias);
+                return res;
+            }
+        }
+
         auto labelOp = builder.create<mlir_ts::LabelOp>(location, builder.getStringAttr(label));
 
         // add merge block
@@ -1041,6 +1076,7 @@ namespace mlirgen
             emitError(location, "A 'continue' statement can only jump to a label of an enclosing statement");
             return mlir::failure();
         }
+        label = resolveLoopLabelAlias(label);
 
         EXIT_IF_FAILED(mlirGenScopeExit(location, DisposeDepth::LoopScope, label, &genContext));
 
@@ -1058,6 +1094,7 @@ namespace mlirgen
             emitError(location, "A 'break' statement can only jump to a label of an enclosing statement");
             return mlir::failure();
         }
+        label = resolveLoopLabelAlias(label);
 
         EXIT_IF_FAILED(mlirGenScopeExit(location, DisposeDepth::LoopScope, label, &genContext));
 

tslang/test/tester/CMakeLists.txt

@@ -332,6 +332,7 @@ tslang_add_test(NAME test-compile-00-while COMMAND test-runner "${PROJECT_SOURCE
 tslang_add_test(NAME test-compile-00-for COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00for.ts")
 tslang_add_test(NAME test-compile-00-for-infinite-result COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_infinite_result.ts")
 tslang_add_test(NAME test-compile-00-break-continue COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00break_continue.ts")
+tslang_add_test(NAME test-compile-00-label-chain-loop COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00label_chain_loop.ts")
 tslang_add_test(NAME test-compile-00-vars COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00vars.ts")
 tslang_add_test(NAME test-compile-00-var-bindings COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00var_bindings.ts")
 tslang_add_test(NAME test-compile-00-globals COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00globals.ts")
@@ -870,6 +871,7 @@ tslang_add_test(NAME test-jit-00-while COMMAND test-runner -jit "${PROJECT_SOURC
 tslang_add_test(NAME test-jit-00-for COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for.ts")
 tslang_add_test(NAME test-jit-00-for-infinite-result COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_infinite_result.ts")
 tslang_add_test(NAME test-jit-00-break-continue COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00break_continue.ts")
+tslang_add_test(NAME test-jit-00-label-chain-loop COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00label_chain_loop.ts")
 tslang_add_test(NAME test-jit-00-vars COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00vars.ts")
 tslang_add_test(NAME test-jit-00-var-bindings COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00var_bindings.ts")
 tslang_add_test(NAME test-jit-00-globals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00globals.ts")
@@ -1819,6 +1821,7 @@ set(TSLANG_CORPUS
     00bool_const_array.ts
     00break_continue_scope_exit.ts
     00break_continue.ts
+    00label_chain_loop.ts
     00capture_in_new_arguments.ts
     00cast_to_f32.ts
     00catch_untyped_value.ts

tslang/test/tester/tests/00label_chain_loop.ts

@@ -0,0 +1,136 @@
+// `a: b: loop`: both labels name the loop. `continue a` goes on to the loop's next iteration, as
+// `continue b` does - it used to restart the loop from its initializer, skipping the incrementor.
+
+function for_continue_outer_label() {
+    let s = "";
+    a: b: for (let i = 0; i < 3; i++) {
+        if (i == 1) continue a;
+        s += i;
+    }
+
+    return s;
+}
+
+function for_continue_inner_label() {
+    let s = "";
+    a: b: for (let i = 0; i < 3; i++) {
+        if (i == 1) continue b;
+        s += i;
+    }
+
+    return s;
+}
+
+function for_continue_from_nested_loop() {
+    let s = "";
+    a: b: c: for (let i = 0; i < 3; i++) {
+        for (let j = 0; j < 3; j++) {
+            if (j == 1) continue a;
+            if (i == 2) break b;
+            s += i;
+            s += j;
+        }
+    }
+
+    return s;
+}
+
+function while_continue_outer_label() {
+    let s = "";
+    let i = 0;
+    a: b: while (i < 4) {
+        i++;
+        if (i == 2) continue a;
+        if (i == 4) break a;
+        s += i;
+    }
+
+    return s;
+}
+
+function do_continue_outer_label() {
+    let s = "";
+    let i = 0;
+    a: b: do {
+        i++;
+        if (i == 2) continue a;
+        s += i;
+    } while (i < 3);
+
+    return s;
+}
+
+function for_of_continue_outer_label() {
+    let s = "";
+    a: b: for (const v of [1, 2, 3]) {
+        if (v == 2) continue a;
+        s += v;
+    }
+
+    return s;
+}
+
+// generating `[1, 2, 3]`'s iterator generated loops of its own, which took the label before the
+// `for` the loop is made into existed: `continue b` then found no loop and crashed the compiler
+function for_of_one_label() {
+    let s = "";
+    b: for (const v of [1, 2, 3]) {
+        if (v == 2) continue b;
+        s += v;
+    }
+
+    const arr: number[] = [4, 5, 6];
+    c: for (const v of arr) {
+        if (v == 5) continue c;
+        if (v == 6) break c;
+        s += v;
+    }
+
+    return s;
+}
+
+// a function inside the loop has labels of its own: its `a` is not another name for the outer `b`
+function nested_function_label() {
+    let s = "";
+    a: b: for (let i = 0; i < 2; i++) {
+        const f = () => {
+            let r = "";
+            a: for (let j = 0; j < 3; j++) {
+                if (j == 1) continue a;
+                r += j;
+            }
+
+            return r;
+        };
+
+        s += f();
+    }
+
+    return s;
+}
+
+function block_chain() {
+    // a chain that ends in a block, not a loop, keeps its labels. A number, not a string: under rc a
+    // jump out of a labelled block releases the function's owned locals (a separate, open bug)
+    let n = 0;
+    a: b: {
+        n++;
+        if (n < 3) continue a;
+    }
+
+    return n;
+}
+
+function main() {
+    assert(for_continue_outer_label() == "02", "for, continue a");
+    assert(for_continue_inner_label() == "02", "for, continue b");
+    assert(for_continue_from_nested_loop() == "0010", "for, continue a from a nested loop");
+    assert(while_continue_outer_label() == "13", "while, continue a");
+    assert(do_continue_outer_label() == "13", "do, continue a");
+    assert(for_of_continue_outer_label() == "13", "for of, continue a");
+    assert(for_of_one_label() == "134", "for of, one label");
+    assert(nested_function_label() == "0202", "function in the loop");
+    assert(block_chain() == 3, "block");
+
+    print("done.");
+}