diff --git a/codon/parser/visitors/typecheck/error.cpp b/codon/parser/visitors/typecheck/error.cpp index 72f7f88bd..18618fb21 100644 --- a/codon/parser/visitors/typecheck/error.cpp +++ b/codon/parser/visitors/typecheck/error.cpp @@ -224,25 +224,33 @@ void TypecheckVisitor::visit(WithStmt *stmt) { std::vector content; for (auto i = stmt->items.size(); i-- > 0;) { - std::string var = stmt->vars[i].empty() ? getTemporaryVar("with") : stmt->vars[i]; - auto as = N(N(var), (*stmt)[i], nullptr, - (*stmt)[i]->hasAttribute(Attr::ExprDominated) - ? AssignStmt::UpdateMode::Update - : AssignStmt::UpdateMode::Assign); + std::string manager = getTemporaryVar("with"); + std::string var = stmt->vars[i]; + auto managerAssign = N(N(manager), (*stmt)[i], nullptr, + AssignStmt::UpdateMode::Assign); Expr *enter = - N(N(N(var), isAsync ? "__aenter__" : "__enter__")); + N(N(N(manager), isAsync ? "__aenter__" : "__enter__")); Expr *exit = - N(N(N(var), isAsync ? "__aexit__" : "__exit__")); + N(N(N(manager), isAsync ? "__aexit__" : "__exit__")); if (isAsync) { enter = N(enter); exit = N(exit); } - content = std::vector{ - as, N(enter), + std::vector wrapped{managerAssign}; + if (var.empty()) { + wrapped.push_back(N(enter)); + } else { + wrapped.push_back(N( + N(var), enter, nullptr, + (*stmt)[i]->hasAttribute(Attr::ExprDominated) ? AssignStmt::UpdateMode::Update + : AssignStmt::UpdateMode::Assign)); + } + wrapped.push_back( N(!content.empty() ? N(content) : clone(stmt->getSuite()), std::vector{}, nullptr, - N(N(exit)))}; + N(N(exit)))); + content = std::move(wrapped); } resultStmt = transform(N(content)); }