diff --git a/integration/tests/for.py b/integration/tests/for.py new file mode 100644 index 0000000..f496ad5 --- /dev/null +++ b/integration/tests/for.py @@ -0,0 +1,10 @@ +def test_else_branch(): + a = [] + for el in ["a", "b", "c"]: + for _ in a: + pass + else: + a.append(el) + return a + +assert test_else_branch() == ["a", "b", "c"] \ No newline at end of file diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index 7de4b40..46373a8 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -1272,6 +1272,11 @@ namespace py { mlir::Block *end_block) const { return [this, &rewriter, condition_start, end_block](mlir::Operation *operation) { + auto parent_is_orelse = [](mlir::Operation *operation) { + auto forloop_op = operation->getParentOfType(); + if (!forloop_op) { return false; } + return &forloop_op.getOrelse() == operation->getParentRegion(); + }; // llvm::outs() << "ForOpLowering 1:\n"; // operation->print(llvm::outs()); // llvm::outs() << '\n'; @@ -1308,6 +1313,9 @@ namespace py { && mlir::isa(yield_op->getParentOp())) { return WalkResult::advance(); } + // is this hacky? maybe there is a better way of ignoring the else branch of + // a for loop + if (parent_is_orelse(operation)) { return WalkResult::advance(); } rewriter.setInsertionPoint(yield_op); if (!yield_op.getKind().has_value() || yield_op.getKind().value() == py::LoopOpKind::continue_) { @@ -2215,6 +2223,11 @@ namespace py { config.enableRegionSimplification = GreedySimplifyRegionLevel::Disabled; config.useTopDownTraversal = true; FrozenRewritePatternSet frozen_patterns{ std::move(patterns) }; + + // getOperation()->print(llvm::outs()); + // llvm::outs() << "-----------------------------------------------\n\n\n"; + // llvm::outs().flush(); + // Currently ignoring the return value as it seems to always fail, even though the // transformation seems to generate the expected output (void)applyPatternsAndFoldGreedily(getOperation(), frozen_patterns, config); diff --git a/src/runtime/PyObject.cpp b/src/runtime/PyObject.cpp index a94eddc..0eeadb4 100644 --- a/src/runtime/PyObject.cpp +++ b/src/runtime/PyObject.cpp @@ -1161,7 +1161,7 @@ PyResult PyObject::__ne__(const PyObject *other) const if (!type_prototype().__eq__.has_value()) { return Ok(not_implemented()); } return call_slot(*type_prototype().__eq__, this, other).and_then([](PyObject *obj) { return truthy(obj, VirtualMachine::the().interpreter()).and_then([](bool value) { - return Ok(value ? py_true() : py_false()); + return Ok(value ? py_false() : py_true()); }); }); }