diff --git a/codon/parser/visitors/simplify/assign.cpp b/codon/parser/visitors/simplify/assign.cpp index 62a4cedf..f5359868 100644 --- a/codon/parser/visitors/simplify/assign.cpp +++ b/codon/parser/visitors/simplify/assign.cpp @@ -18,6 +18,8 @@ namespace codon::ast { void SimplifyVisitor::visit(AssignExpr *expr) { seqassert(expr->var->getId(), "only simple assignment expression are supported"); StmtPtr s = N(clone(expr->var), expr->expr); + auto avoidDomination = false; // walruses always leak + std::swap(avoidDomination, ctx->avoidDomination); if (ctx->isConditionalExpr) { // Make sure to transform both suite _AND_ the expression in the same scope ctx->enterConditionalBlock(); @@ -33,6 +35,7 @@ void SimplifyVisitor::visit(AssignExpr *expr) { s = transform(s); transform(expr->var); } + std::swap(avoidDomination, ctx->avoidDomination); resultExpr = N(std::vector{s}, expr->var); } @@ -155,7 +158,11 @@ StmtPtr SimplifyVisitor::transformAssignment(ExprPtr lhs, ExprPtr rhs, ExprPtr t val = ctx->addVar(e->value, canonical, lhs->getSrcInfo()); if (auto st = getStaticGeneric(type.get())) val->staticType = st; + if (ctx->avoidDomination) + val->avoidDomination = true; } + // Clean up seen tags if shadowing a name + ctx->seenGlobalIdentifiers[ctx->getBaseName()].erase(e->value); // Register all toplevel variables as global in JIT mode bool isGlobal = (ctx->cache->isJit && val->isGlobal() && !val->isGeneric()) || diff --git a/codon/parser/visitors/simplify/collections.cpp b/codon/parser/visitors/simplify/collections.cpp index 21e561ab..9c4f7c4e 100644 --- a/codon/parser/visitors/simplify/collections.cpp +++ b/codon/parser/visitors/simplify/collections.cpp @@ -65,6 +65,8 @@ void SimplifyVisitor::visit(GeneratorExpr *expr) { } SuiteStmt *prev = nullptr; + auto avoidDomination = true; + std::swap(avoidDomination, ctx->avoidDomination); auto suite = transformGeneratorBody(loops, prev); ExprPtr var = N(ctx->cache->getTemporaryVar("gen")); if (expr->kind == GeneratorExpr::ListGenerator) { @@ -94,6 +96,7 @@ void SimplifyVisitor::visit(GeneratorExpr *expr) { stmts.push_back(suite); resultExpr = N(N(N(makeAnonFn(stmts)), "__iter__")); } + std::swap(avoidDomination, ctx->avoidDomination); } /// Transform a dictionary comprehension to the corresponding statement expression. @@ -102,6 +105,8 @@ void SimplifyVisitor::visit(GeneratorExpr *expr) { /// for i in j: if a: gen.__setitem__(i+a, j+1)``` void SimplifyVisitor::visit(DictGeneratorExpr *expr) { SuiteStmt *prev = nullptr; + auto avoidDomination = true; + std::swap(avoidDomination, ctx->avoidDomination); auto suite = transformGeneratorBody(expr->loops, prev); std::vector stmts; @@ -111,6 +116,7 @@ void SimplifyVisitor::visit(DictGeneratorExpr *expr) { clone(expr->key), clone(expr->expr)))); stmts.push_back(transform(suite)); resultExpr = N(stmts, transform(var)); + std::swap(avoidDomination, ctx->avoidDomination); } /// Transforms a list of @c GeneratorBody loops to the corresponding set of for loops. @@ -127,7 +133,8 @@ StmtPtr SimplifyVisitor::transformGeneratorBody(const std::vector newSuite = N(); auto nextPrev = dynamic_cast(newSuite.get()); - prev->stmts.push_back(N(l.vars->clone(), l.gen->clone(), newSuite)); + auto forStmt = N(l.vars->clone(), l.gen->clone(), newSuite); + prev->stmts.push_back(forStmt); prev = nextPrev; for (auto &cond : l.conds) { newSuite = N(); diff --git a/codon/parser/visitors/simplify/ctx.cpp b/codon/parser/visitors/simplify/ctx.cpp index 54c3c643..57c8f2eb 100644 --- a/codon/parser/visitors/simplify/ctx.cpp +++ b/codon/parser/visitors/simplify/ctx.cpp @@ -171,6 +171,8 @@ SimplifyContext::Item SimplifyContext::findDominatingBinding(const std::string & for (auto i = it->second.begin(); i != it->second.end(); i++) { if (i == lastGood) break; + if (!(*i)->canDominate()) + continue; // These bindings (and their canonical identifiers) will be replaced by the // dominating binding during the type checking pass. cache->replacements[(*i)->canonicalName] = {canonicalName, hasUsed}; diff --git a/codon/parser/visitors/simplify/ctx.h b/codon/parser/visitors/simplify/ctx.h index d17272bf..4e5c6a4e 100644 --- a/codon/parser/visitors/simplify/ctx.h +++ b/codon/parser/visitors/simplify/ctx.h @@ -44,6 +44,9 @@ struct SimplifyItem : public SrcObject { bool generic = false; /// Set if an identifier is a static variable. char staticType = 0; + /// Set if an identifier should not be dominated + /// (e.g., a loop variable in a comprehension). + bool avoidDomination = false; public: SimplifyItem(Kind kind, std::string baseName, std::string canonicalName, @@ -66,6 +69,8 @@ struct SimplifyItem : public SrcObject { bool isConditional() const { return scope.size() > 1; } bool isGeneric() const { return generic; } char isStatic() const { return staticType; } + /// True if an identifier is a loop variable in a comprehension + bool canDominate() const { return !avoidDomination; } }; /** Context class that tracks identifiers during the simplification. **/ @@ -159,6 +164,8 @@ struct SimplifyContext : public Context { /// Allow type() expressions. Currently used to disallow type() in class /// and function definitions. bool allowTypeOf; + /// Set if all assignments should not be dominated later on. + bool avoidDomination = false; public: SimplifyContext(std::string filename, Cache *cache); diff --git a/codon/parser/visitors/simplify/loops.cpp b/codon/parser/visitors/simplify/loops.cpp index e0b4f3c2..eee3fc96 100644 --- a/codon/parser/visitors/simplify/loops.cpp +++ b/codon/parser/visitors/simplify/loops.cpp @@ -98,13 +98,14 @@ void SimplifyVisitor::visit(ForStmt *stmt) { ctx->getBase()->loops.push_back({breakVar, ctx->scope.blocks, {}}); std::string varName; if (auto i = stmt->var->getId()) { - ctx->addVar(i->value, varName = ctx->generateCanonicalName(i->value), - stmt->var->getSrcInfo()); + auto val = ctx->addVar(i->value, varName = ctx->generateCanonicalName(i->value), + stmt->var->getSrcInfo()); + val->avoidDomination = ctx->avoidDomination; transform(stmt->var); stmt->suite = transform(N(stmt->suite)); } else { varName = ctx->cache->getTemporaryVar("for"); - ctx->addVar(varName, varName, stmt->var->getSrcInfo()); + auto val = ctx->addVar(varName, varName, stmt->var->getSrcInfo()); auto var = N(varName); std::vector stmts; // Add for_var = [for variables] diff --git a/codon/parser/visitors/typecheck/infer.cpp b/codon/parser/visitors/typecheck/infer.cpp index 601c8556..62ab3115 100644 --- a/codon/parser/visitors/typecheck/infer.cpp +++ b/codon/parser/visitors/typecheck/infer.cpp @@ -325,6 +325,7 @@ types::TypePtr TypecheckVisitor::realizeFunc(types::FuncType *type, bool force) // Lambda typecheck failures are "ignored" as they are treated as statements, // not functions. // TODO: generalize this further. + // LOG("{}", ast->suite->toString(2)); error("cannot typecheck the program"); } ctx->realizationBases.pop_back(); diff --git a/test/parser/simplify_stmt.codon b/test/parser/simplify_stmt.codon index 98aedf87..faeabbfa 100644 --- a/test/parser/simplify_stmt.codon +++ b/test/parser/simplify_stmt.codon @@ -1209,6 +1209,19 @@ for i in range(2): #: 1 #: 1 +def comprehension_test(x): + for n in range(3): + print('>', n) + l = ['1', '2', str(x)] + x = [n for n in l] + print(x, n) +comprehension_test(5) +#: > 0 +#: > 1 +#: > 2 +#: ['1', '2', '5'] 2 + + #%% block_unroll,barebones # Ensure that block unrolling is done in RAII manner on error def foo():