mlir: bind loop-else break/continue transitively

Author: gf712

integration/tests/loop_else_break_binding.py

@@ -13,9 +13,13 @@
     producing a cross-region block reference the verifier rejects.
 
 Both are now handled by deferring: a loop refuses to lower while a nested loop
-still holds a break/continue in its orelse, so the nested loop is flattened into
-the enclosing region first and the branch is same-region by construction. That
-handshake is also why both loop patterns share one pass.
+still holds a break/continue that binds to it, so the nested loop is flattened
+into the enclosing region first and the branch is same-region by construction.
+That handshake is also why both loop patterns share one pass.
+
+Binding outwards is transitive: an else nested inside another else is still
+lexically part of whichever loop body encloses the pair, so the deferral has to
+follow the whole chain rather than stop one level in.
 """
 
 # break in a nested for's else breaks the OUTER for.
@@ -107,3 +111,76 @@
         log.append(("else", outer))
     log.append(("after", outer))
 assert log == [("else", 1), ("after", 1), ("else", 2), ("after", 2)], log
+
+# An else nested inside another else: the break binds outwards through *both*, to
+# the outermost loop. Checking only the first nested loop's else missed this and
+# silently dropped the break.
+log = []
+for a in [1, 2, 3]:
+    log.append(a)
+    for b in []:
+        pass
+    else:
+        for c in []:
+            pass
+        else:
+            break
+assert log == [1], log
+
+# Same shape with `continue`, which must skip the rest of the outermost body.
+log = []
+for a in [1, 2, 3]:
+    log.append(a)
+    for b in []:
+        pass
+    else:
+        for c in []:
+            pass
+        else:
+            continue
+    log.append("after-must-not-run")
+assert log == [1, 2, 3], log
+
+# Same shape built from `while`s. Here the mis-binding was not silent: the branch
+# was emitted while the inner py.while was still a region of its own, which the
+# verifier rejects as a reference to a block in another region.
+log = []
+for a in [1, 2, 3]:
+    log.append(a)
+    while False:
+        pass
+    else:
+        while False:
+            pass
+        else:
+            break
+assert log == [1], log
+
+# Three elses deep, to check the walk follows the chain rather than a fixed depth.
+log = []
+for a in [1, 2, 3]:
+    log.append(a)
+    for b in []:
+        pass
+    else:
+        for c in []:
+            pass
+        else:
+            for d in []:
+                pass
+            else:
+                break
+assert log == [1], log
+
+# The chain stops at the first loop *body*: this break is in the body of a loop
+# that happens to sit in an else, so it binds to that loop and no further.
+log = []
+for a in [1, 2]:
+    for b in []:
+        pass
+    else:
+        for c in [10, 20]:
+            log.append((a, c))
+            break
+    log.append(("after", a))
+assert log == [(1, 10), ("after", 1), (2, 10), ("after", 2)], log

src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp

@@ -56,39 +56,54 @@ namespace py {
 			return mlir::isa<TryOp, ForLoopOp, WithOp, WhileOp, TryHandlerOp>(op);
 		}
 
-		// True when `yield_op` is a loop-control (break/continue) yield that binds to
-		// the loop *enclosing* `loop` rather than to `loop` itself — i.e. it sits in
-		// `loop`'s orelse, which is not part of the loop body.
-		bool binds_to_enclosing_loop(mlir::py::PyLoopOpInterface loop,
-			mlir::py::BranchYieldOp yield_op)
+		// True when `yield_op`, a loop-control (break/continue) yield, binds to the
+		// loop whose body region is `body`.
+		//
+		// Python binds break/continue to the innermost loop whose *body* lexically
+		// contains it. An else clause is not part of its own loop's body, so a yield
+		// sitting there keeps searching outwards — and transitively so: an else nested
+		// inside another else is still lexically part of whatever body encloses the
+		// pair. Regions that are neither body nor orelse (a try body, a with body) are
+		// likewise transparent, which is what makes `break` inside a `try` bind to the
+		// loop around it.
+		bool binds_to_loop(mlir::Region &body, mlir::py::BranchYieldOp yield_op)
 		{
-			return yield_op.getKind().has_value() && loop.isLoopOrelse(yield_op->getParentRegion());
+			for (mlir::Region *region = yield_op->getParentRegion(); region != nullptr;
+				region = region->getParentRegion()) {
+				if (region == &body) { return true; }
+				auto loop =
+					mlir::dyn_cast_if_present<mlir::py::PyLoopOpInterface>(region->getParentOp());
+				// A loop body (or a for's step) stops the search: the yield is that
+				// loop's, and its own pattern claims it.
+				if (loop && !loop.isLoopOrelse(region)) { return false; }
+			}
+			return false;
 		}
 
-		// True when some loop nested in `region` still holds a break/continue that
-		// binds to the loop being lowered — i.e. one sitting in that nested loop's
-		// orelse. Such a yield cannot be rewritten yet: it lives in a region that has
-		// not been flattened, so branching it to our target block would be a
-		// cross-region block reference, which is invalid IR.
+		// True when a break/continue that binds to the loop whose body is `body` is
+		// somewhere replace_loop_branch_yields cannot reach yet — inside a nested
+		// region that has not been flattened into ours. Branching it to our target
+		// block now would be a cross-region block reference, which is invalid IR.
 		//
-		// The caller defers (fails the match) until the nested loop lowers and inlines
+		// The caller defers (fails the match) until the nested op lowers and inlines
 		// the yield into our region, the same innermost-first trick TryOpLowering uses
-		// for nested trys. Terminates because the innermost such loop has nothing
-		// nested to wait on.
-		bool has_pending_nested_orelse_control(mlir::Region &region)
+		// for nested trys. Terminates because the innermost such op has nothing nested
+		// to wait on.
+		bool has_pending_nested_orelse_control(mlir::Region &body)
 		{
-			if (region.empty()) { return false; }
+			if (body.empty()) { return false; }
 			bool pending = false;
-			region.walk<WalkOrder::PreOrder>([&pending](mlir::Operation *op) {
-				auto loop = mlir::dyn_cast<mlir::py::PyLoopOpInterface>(op);
-				if (!loop) { return WalkResult::advance(); }
-				loop.getLoopOrelseRegion().walk<WalkOrder::PreOrder>(
-					[&pending, loop](mlir::py::BranchYieldOp yield_op) {
-						if (binds_to_enclosing_loop(loop, yield_op)) { pending = true; }
-					});
-				// Only this loop's own orelse matters here; anything deeper is the
-				// nested loop's problem and it defers on it in turn.
-				return WalkResult::skip();
+			body.walk<WalkOrder::PreOrder>([&pending, &body](mlir::Operation *op) {
+				// Mirror replace_loop_branch_yields: what it walks through it rewrites
+				// in place, so only what it skips over can be pending.
+				if (!is_flattened_region_op(op)) { return WalkResult::advance(); }
+				op->walk([&pending, &body](mlir::py::BranchYieldOp yield_op) {
+					if (!yield_op.getKind().has_value()) { return WalkResult::advance(); }
+					if (!binds_to_loop(body, yield_op)) { return WalkResult::advance(); }
+					pending = true;
+					return WalkResult::interrupt();
+				});
+				return pending ? WalkResult::interrupt() : WalkResult::skip();
 			});
 			return pending;
 		}