An if/else branch written without braces is a scope of its own (#458)

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -4239,6 +4239,7 @@ class MLIRGenImpl
     }
 
     mlir::LogicalResult mlirGen(IfStatement ifStatementAST, const GenContext &genContext);
+    mlir::LogicalResult mlirGenBranchStatement(Statement statement, const GenContext &genContext);
 
     mlir::LogicalResult mlirGen(DoStatement doStatementAST, const GenContext &genContext);
 

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -611,7 +611,7 @@ namespace mlirgen
                 // check if we do safe-cast here
                 checkSafeCast(ifStatementAST->expression, V(result), hasElse || thenExits ? &elseSafeCase : nullptr, genContext);
 
-                auto result = mlirGen(ifStatementAST->thenStatement, genContext);
+                auto result = mlirGenBranchStatement(ifStatementAST->thenStatement, genContext);
                 EXIT_IF_FAILED(result)
             }
         }
@@ -632,7 +632,7 @@ namespace mlirgen
                     addSafeCastStatement(elseSafeCase.expr, elseSafeCase.safeType, false, nullptr, genContext);
                 }
 
-                auto result = mlirGen(ifStatementAST->elseStatement, genContext);
+                auto result = mlirGenBranchStatement(ifStatementAST->elseStatement, genContext);
                 EXIT_IF_FAILED(result)
             }
         }
@@ -647,6 +647,35 @@ namespace mlirgen
         return mlir::success();
     }
 
+    // A branch written without braces (`else for (const c of gen()) ...`) is a scope all the same:
+    // what it registers to give back at the end of its block - an owned local such as a
+    // generator's iterator, a `using` - belongs to the branch, whose region it is defined in, not
+    // to the block around the `if`, whose end is outside that region (#458). A braced branch is a
+    // Block and makes this scope itself.
+    mlir::LogicalResult MLIRGenImpl::mlirGenBranchStatement(Statement statement, const GenContext &genContext)
+    {
+        if (statement == SyntaxKind::Block)
+        {
+            return mlirGen(statement, genContext);
+        }
+
+        auto location = loc(statement);
+
+        GenContext branchContext(genContext);
+        branchContext.parentBlockContext = &genContext;
+        branchContext.isLoopBodyScope = false;
+
+        auto usingVars = std::make_unique<SmallVector<ts::VariableDeclarationDOM::TypePtr>>();
+        branchContext.usingVars = usingVars.get();
+
+        auto ownedVars = std::make_unique<SmallVector<mlir::Value>>();
+        branchContext.ownedVars = ownedVars.get();
+
+        EXIT_IF_FAILED(mlirGen(statement, branchContext));
+
+        return mlirGenScopeExit(location, DisposeDepth::CurrentScope, {}, &branchContext);
+    }
+
     mlir::LogicalResult MLIRGenImpl::mlirGen(DoStatement doStatementAST, const GenContext &genContext)
     {
         SymbolTableScopeT varScope(symbolTable);

tslang/test/tester/CMakeLists.txt

@@ -471,6 +471,7 @@ tslang_add_test(NAME test-compile-00-class-structural-extends COMMAND test-runne
 tslang_add_test(NAME test-compile-00-global-constant-concat COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_constant_concat.ts")
 tslang_add_test(NAME test-compile-00-global-object-literal COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_object_literal.ts")
 tslang_add_test(NAME test-compile-00-instanceof COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof.ts")
+tslang_add_test(NAME test-compile-00-unbraced-branch-scope COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00unbraced_branch_scope.ts")
 tslang_add_test(NAME test-compile-00-const-of-let COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_of_let.ts")
 tslang_add_test(NAME test-compile-00-instanceof-union COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof_union.ts")
 tslang_add_test(NAME test-compile-00-const-object-literal-field COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_object_literal_field.ts")
@@ -963,6 +964,7 @@ tslang_add_test(NAME test-jit-00-class-structural-extends COMMAND test-runner -j
 tslang_add_test(NAME test-jit-00-global-constant-concat COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_constant_concat.ts")
 tslang_add_test(NAME test-jit-00-global-object-literal COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_object_literal.ts")
 tslang_add_test(NAME test-jit-00-instanceof COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof.ts")
+tslang_add_test(NAME test-jit-00-unbraced-branch-scope COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00unbraced_branch_scope.ts")
 tslang_add_test(NAME test-jit-00-const-of-let COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_of_let.ts")
 tslang_add_test(NAME test-jit-00-instanceof-union COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof_union.ts")
 tslang_add_test(NAME test-jit-00-const-object-literal-field COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_object_literal_field.ts")
@@ -1787,6 +1789,7 @@ set(TSLANG_CORPUS
     00const_object_literal_field.ts
     00global_union_from_call.ts
     00instanceof.ts
+    00unbraced_branch_scope.ts
     00global_constant_concat.ts
     00const_of_let.ts
     00instanceof_namespace_narrowing.ts

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

@@ -0,0 +1,36 @@
+// A branch written without braces is a scope of its own (#458): what it declares and must give back
+// at the end - the iterator of a `for ... of` over a generator, a `using` - is released at the end
+// of the branch. It was handed to the block around the `if`, whose end is outside the branch, and
+// the IR failed to verify: "operand #0 does not dominate this use".
+function* gen() {
+    yield 1;
+    yield 2;
+}
+
+class Res {
+    [Symbol.dispose]() {
+        print("disposed");
+    }
+}
+
+function main() {
+    let sum = 0;
+    const it = gen();
+    if (false) it.next(); else for (const c of it) sum += c;
+    assert(sum == 3, "an unbraced else that iterates a generator");
+
+    for (let i = 0; i < 2; i++) {
+        if (i == 0) sum += 10; else for (const c of gen()) sum += c;
+    }
+
+    assert(sum == 16, "the same inside a loop");
+
+    if (sum > 0) for (const c of gen()) sum += c;
+    assert(sum == 19, "an unbraced then");
+
+    let total = 0;
+    for (let i = 0; i < 3; i++) if (i != 1) for (const c of gen()) total += c;
+    assert(total == 6, "an unbraced if as a loop body");
+
+    print("done.");
+}