[TypeChecker] Split for-in sequence into parsed and type-checked versions

xedin · xedin · commit b7860ea05531 · 2022-05-30T23:17:41.000-07:00
diff --git a/include/swift/AST/Stmt.h b/include/swift/AST/Stmt.h
@@ -743,7 +743,7 @@ class ForEachStmt : public LabeledStmt {
 
   // Set by Sema:
   ProtocolConformanceRef sequenceConformance = ProtocolConformanceRef();
-  VarDecl *iteratorVar = nullptr;
+  PatternBindingDecl *iteratorVar = nullptr;
   Expr *nextCall = nullptr;
   OpaqueValueExpr *elementExpr = nullptr;
   Expr *convertElementExpr = nullptr;
@@ -759,8 +759,8 @@ class ForEachStmt : public LabeledStmt {
     setPattern(Pat);
   }
 
-  void setIteratorVar(VarDecl *var) { iteratorVar = var; }
-  VarDecl *getIteratorVar() const { return iteratorVar; }
+  void setIteratorVar(PatternBindingDecl *var) { iteratorVar = var; }
+  PatternBindingDecl *getIteratorVar() const { return iteratorVar; }
 
   void setNextCall(Expr *next) { nextCall = next; }
   Expr *getNextCall() const { return nextCall; }
@@ -802,8 +802,12 @@ class ForEachStmt : public LabeledStmt {
   /// by this foreach loop, as it was written in the source code and
   /// subsequently type-checked. To determine the semantic behavior of this
   /// expression to extract a range, use \c getRangeInit().
-  Expr *getSequence() const { return Sequence; }
-  void setSequence(Expr *S) { Sequence = S; }
+  Expr *getParsedSequence() const { return Sequence; }
+  void setParsedSequence(Expr *S) { Sequence = S; }
+
+  /// Type-checked version of the sequence or nullptr if this statement
+  /// yet to be type-checked.
+  Expr *getTypeCheckedSequence() const;
 
   /// getBody - Retrieve the body of the loop.
   BraceStmt *getBody() const { return Body; }
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -1011,7 +1011,7 @@ struct ForEachStmtInfo {
   Type initType;
 
   /// Implicit `$iterator = <sequence>.makeIterator()`
-  VarDecl *makeIteratorVar;
+  PatternBindingDecl *makeIteratorVar;
 
   /// Implicit `$iterator.next()` call.
   Expr *nextCall;
@@ -2434,7 +2434,7 @@ class SolutionApplicationTarget {
     case Kind::forEachStmt:
       auto *stmt = forEachStmt.stmt;
       SourceLoc startLoc = stmt->getForLoc();
-      SourceLoc endLoc = stmt->getSequence()->getEndLoc();
+      SourceLoc endLoc = stmt->getParsedSequence()->getEndLoc();
 
       if (auto *whereExpr = stmt->getWhere()) {
         endLoc = whereExpr->getEndLoc();
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -1572,7 +1572,7 @@ class PrintStmt : public StmtVisitor<PrintStmt> {
     }
     printRec(S->getPattern());
     OS << '\n';
-    printRec(S->getSequence());
+    printRec(S->getParsedSequence());
     OS << '\n';
     if (S->getIteratorVar()) {
       printRec(S->getIteratorVar());
diff --git a/lib/AST/ASTScopeCreation.cpp b/lib/AST/ASTScopeCreation.cpp
@@ -994,7 +994,7 @@ void SwitchStmtScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
 
 void ForEachStmtScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
     ScopeCreator &scopeCreator) {
-  scopeCreator.addToScopeTree(stmt->getSequence(), this);
+  scopeCreator.addToScopeTree(stmt->getParsedSequence(), this);
 
   // Add a child describing the scope of the pattern.
   // In error cases such as:
diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -1590,11 +1590,28 @@ Stmt *Traversal::visitForEachStmt(ForEachStmt *S) {
 
   // The iterator decl is built directly on top of the sequence
   // expression, so don't visit both.
-  if (Expr *Sequence = S->getSequence()) {
-    if ((Sequence = doIt(Sequence)))
-      S->setSequence(Sequence);
-    else
-      return nullptr;
+  //
+  // If for-in is already type-checked, the type-checked version
+  // of the sequence is going to be visited as part of `iteratorVar`.
+  if (S->getTypeCheckedSequence()) {
+    if (auto IteratorVar = S->getIteratorVar()) {
+      if (doIt(IteratorVar))
+        return nullptr;
+    }
+
+    if (auto NextCall = S->getNextCall()) {
+      if ((NextCall = doIt(NextCall)))
+        S->setNextCall(NextCall);
+      else
+        return nullptr;
+    }
+  } else {
+    if (Expr *Sequence = S->getParsedSequence()) {
+      if ((Sequence = doIt(Sequence)))
+        S->setParsedSequence(Sequence);
+      else
+        return nullptr;
+    }
   }
 
   if (Expr *Where = S->getWhere()) {
@@ -1611,18 +1628,6 @@ Stmt *Traversal::visitForEachStmt(ForEachStmt *S) {
       return nullptr;
   }
 
-  if (auto IteratorVar = S->getIteratorVar()) {
-    if (doIt(IteratorVar))
-      return nullptr;
-  }
-
-  if (auto NextCall = S->getNextCall()) {
-    if ((NextCall = doIt(NextCall)))
-      S->setNextCall(NextCall);
-    else
-      return nullptr;
-  }
-
   if (Stmt *Body = S->getBody()) {
     if ((Body = doIt(Body)))
       S->setBody(cast<BraceStmt>(Body));
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -3110,7 +3110,7 @@ void FindLocalVal::visitForEachStmt(ForEachStmt *S) {
   if (!isReferencePointInRange(S->getSourceRange()))
     return;
   visit(S->getBody());
-  if (!isReferencePointInRange(S->getSequence()->getSourceRange()))
+  if (!isReferencePointInRange(S->getParsedSequence()->getSourceRange()))
     checkPattern(S->getPattern(), DeclVisibilityKind::LocalVariable);
 }
 
diff --git a/lib/AST/Stmt.cpp b/lib/AST/Stmt.cpp
@@ -301,6 +301,10 @@ void ForEachStmt::setPattern(Pattern *p) {
   Pat->markOwnedByStatement(this);
 }
 
+Expr *ForEachStmt::getTypeCheckedSequence() const {
+  return iteratorVar ? iteratorVar->getInit(/*index=*/0) : nullptr;
+}
+
 DoCatchStmt *DoCatchStmt::create(ASTContext &ctx, LabeledStmtInfo labelInfo,
                                  SourceLoc doLoc, Stmt *body,
                                  ArrayRef<CaseStmt *> catches,
diff --git a/lib/IDE/ExprContextAnalysis.cpp b/lib/IDE/ExprContextAnalysis.cpp
@@ -1092,7 +1092,7 @@ class ExprContextAnalyzer {
       break;
     }
     case StmtKind::ForEach:
-      if (auto SEQ = cast<ForEachStmt>(Parent)->getSequence()) {
+      if (auto SEQ = cast<ForEachStmt>(Parent)->getParsedSequence()) {
         if (containsTarget(SEQ)) {
           recordPossibleType(Context.getSequenceType());
         }
diff --git a/lib/IDE/Formatting.cpp b/lib/IDE/Formatting.cpp
@@ -2246,7 +2246,7 @@ class FormatWalker : public ASTWalker {
         if (Range.isValid() && overlapsTarget(Range))
           return IndentContext {ForLoc, !OutdentChecker::hasOutdent(SM, P)};
       }
-      if (auto *E = FS->getSequence()) {
+      if (auto *E = FS->getParsedSequence()) {
         SourceRange Range = FS->getInLoc();
         widenOrSet(Range, E->getSourceRange());
         if (Range.isValid() && isTargetContext(Range)) {
diff --git a/lib/IDE/SyntaxModel.cpp b/lib/IDE/SyntaxModel.cpp
@@ -747,8 +747,8 @@ std::pair<bool, Stmt *> ModelASTWalker::walkToStmtPre(Stmt *S) {
                                  charSourceRangeFromSourceRange(SM, ElemRange));
       }
     }
-    if (ForEachS->getSequence())
-      addExprElem(SyntaxStructureElementKind::Expr, ForEachS->getSequence(),SN);
+    if (auto *S = ForEachS->getParsedSequence())
+      addExprElem(SyntaxStructureElementKind::Expr, S, SN);
     pushStructureNode(SN, S);
 
   } else if (auto *WhileS = dyn_cast<WhileStmt>(S)) {
diff --git a/lib/SILGen/SILGenStmt.cpp b/lib/SILGen/SILGenStmt.cpp
@@ -1022,7 +1022,7 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
   // Emit the 'iterator' variable that we'll be using for iteration.
   LexicalScope OuterForScope(SGF, CleanupLocation(S));
   {
-    SGF.emitPatternBinding(S->getIteratorVar()->getParentPatternBinding(),
+    SGF.emitPatternBinding(S->getIteratorVar(),
                            /*index=*/0);
   }
 
@@ -1074,7 +1074,7 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
       SILLocation loc = SILLocation(S);
       RValue result = buildElementRValue(SGFContext(nextInit.get()));
       if (!result.isInContext()) {
-        ArgumentSource(SILLocation(S->getSequence()),
+        ArgumentSource(SILLocation(S->getTypeCheckedSequence()),
                        std::move(result).ensurePlusOne(SGF, loc))
             .forwardInto(SGF, nextInit.get());
       }
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -8556,7 +8556,7 @@ static Optional<SolutionApplicationTarget> applySolutionToForEachStmt(
   auto resultTarget = target;
   auto &forEachStmtInfo = resultTarget.getForEachStmtInfo();
   auto *stmt = target.getAsForEachStmt();
-  auto *parsedSequence = stmt->getSequence();
+  auto *parsedSequence = stmt->getParsedSequence();
   bool isAsync = stmt->getAwaitLoc().isValid();
 
   // Simplify the various types.
@@ -8572,21 +8572,22 @@ static Optional<SolutionApplicationTarget> applySolutionToForEachStmt(
 
   // First, let's apply the solution to the sequence expression.
   {
+    auto *makeIteratorVar = forEachStmtInfo.makeIteratorVar;
+
     auto makeIteratorTarget =
-        *cs.getSolutionApplicationTarget(forEachStmtInfo.makeIteratorVar);
+        *cs.getSolutionApplicationTarget({makeIteratorVar, /*index=*/0});
 
     auto rewrittenTarget = rewriteTarget(makeIteratorTarget);
     if (!rewrittenTarget)
       return None;
 
-    stmt->setSequence(rewrittenTarget->getAsExpr());
-    stmt->setIteratorVar(forEachStmtInfo.makeIteratorVar);
-
-    auto *PB = forEachStmtInfo.makeIteratorVar->getParentPatternBinding();
+    // Set type-checked initializer and mark it as such.
     {
-      PB->setInit(/*index=*/0, stmt->getSequence());
-      PB->setInitializerChecked(/*index=*/0);
+      makeIteratorVar->setInit(/*index=*/0, rewrittenTarget->getAsExpr());
+      makeIteratorVar->setInitializerChecked(/*index=*/0);
     }
+
+    stmt->setIteratorVar(makeIteratorVar);
   }
 
   // Now, `$iterator.next()` call.
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -3886,7 +3886,7 @@ generateForEachStmtConstraints(
   auto forEachStmtInfo = target.getForEachStmtInfo();
   ForEachStmt *stmt = target.getAsForEachStmt();
   bool isAsync = stmt->getAwaitLoc().isValid();
-  auto *sequenceExpr = stmt->getSequence();
+  auto *sequenceExpr = stmt->getParsedSequence();
   auto *dc = target.getDeclContext();
   auto contextualLocator = cs.getConstraintLocator(
       sequenceExpr, LocatorPathElt::ContextualType(CTP_ForEachSequence));
@@ -3940,7 +3940,7 @@ generateForEachStmtConstraints(
                                FreeTypeVariableBinding::Disallow))
       return None;
 
-    forEachStmtInfo.makeIteratorVar = makeIteratorVar;
+    forEachStmtInfo.makeIteratorVar = PB;
 
     // Type of sequence expression has to conform to Sequence protocol.
     {
@@ -3951,7 +3951,7 @@ generateForEachStmtConstraints(
       forEachStmtInfo.sequenceType = cs.getType(sequenceExpr);
     }
 
-    cs.setSolutionApplicationTarget(makeIteratorVar, makeIteratorTarget);
+    cs.setSolutionApplicationTarget({PB, /*index=*/0}, makeIteratorTarget);
   }
 
   // Now, result type of `.makeIterator()` is used to form a call to
diff --git a/lib/Sema/MiscDiagnostics.cpp b/lib/Sema/MiscDiagnostics.cpp
@@ -3784,7 +3784,7 @@ static void checkStmtConditionTrailingClosure(ASTContext &ctx, const Stmt *S) {
   } else if (auto SS = dyn_cast<SwitchStmt>(S)) {
     checkStmtConditionTrailingClosure(ctx, SS->getSubjectExpr());
   } else if (auto FES = dyn_cast<ForEachStmt>(S)) {
-    checkStmtConditionTrailingClosure(ctx, FES->getSequence());
+    checkStmtConditionTrailingClosure(ctx, FES->getParsedSequence());
     checkStmtConditionTrailingClosure(ctx, FES->getWhere());
   } else if (auto DCS = dyn_cast<DoCatchStmt>(S)) {
     for (auto CS : DCS->getCatches())
diff --git a/lib/Sema/PCMacro.cpp b/lib/Sema/PCMacro.cpp
@@ -224,7 +224,7 @@ class Instrumenter : InstrumenterBase {
 
       // point at the for stmt, to look nice
       SourceLoc StartLoc = FES->getStartLoc();
-      SourceLoc EndLoc = FES->getSequence()->getEndLoc();
+      SourceLoc EndLoc = FES->getParsedSequence()->getEndLoc();
       // FIXME: get the 'end' of the for stmt
       // if (FD->getResultTypeRepr()) {
       //   EndLoc = FD->getResultTypeSourceRange().End;
diff --git a/lib/Sema/PreCheckExpr.cpp b/lib/Sema/PreCheckExpr.cpp
@@ -2126,7 +2126,7 @@ bool ConstraintSystem::preCheckTarget(SolutionApplicationTarget &target,
   if (target.isForEachStmt()) {
     auto *stmt = target.getAsForEachStmt();
 
-    auto *sequenceExpr = stmt->getSequence();
+    auto *sequenceExpr = stmt->getParsedSequence();
     auto *whereExpr = stmt->getWhere();
 
     hadErrors |= preCheckExpression(sequenceExpr, DC,
@@ -2141,7 +2141,7 @@ bool ConstraintSystem::preCheckTarget(SolutionApplicationTarget &target,
 
     // Update sequence and where expressions to pre-checked versions.
     if (!hadErrors) {
-      stmt->setSequence(sequenceExpr);
+      stmt->setParsedSequence(sequenceExpr);
 
       if (whereExpr)
         stmt->setWhere(whereExpr);
diff --git a/lib/Sema/TypeCheckConstraints.cpp b/lib/Sema/TypeCheckConstraints.cpp
@@ -306,7 +306,8 @@ void constraints::performSyntacticDiagnosticsForTarget(
     auto *stmt = target.getAsForEachStmt();
 
     // First emit diagnostics for the main expression.
-    performSyntacticExprDiagnostics(stmt->getSequence(), dc, isExprStmt,
+    performSyntacticExprDiagnostics(stmt->getTypeCheckedSequence(), dc,
+                                    isExprStmt,
                                     disableExprAvailabilityChecking);
 
     if (auto *whereExpr = stmt->getWhere())

Original file line number	Diff line number	Diff line change
`@@ -1572,7 +1572,7 @@ class PrintStmt : public StmtVisitor<PrintStmt> {`
`1572`	`1572`	`}`
`1573`	`1573`	`printRec(S->getPattern());`
`1574`	`1574`	`OS << '\n';`
`1575`		`- printRec(S->getSequence());`
	`1575`	`+ printRec(S->getParsedSequence());`
`1576`	`1576`	`OS << '\n';`
`1577`	`1577`	`if (S->getIteratorVar()) {`
`1578`	`1578`	`printRec(S->getIteratorVar());`
Original file line number	Diff line number	Diff line change
`@@ -3110,7 +3110,7 @@ void FindLocalVal::visitForEachStmt(ForEachStmt *S) {`
`3110`	`3110`	`if (!isReferencePointInRange(S->getSourceRange()))`
`3111`	`3111`	`return;`
`3112`	`3112`	`visit(S->getBody());`
`3113`		`- if (!isReferencePointInRange(S->getSequence()->getSourceRange()))`
	`3113`	`+ if (!isReferencePointInRange(S->getParsedSequence()->getSourceRange()))`
`3114`	`3114`	`checkPattern(S->getPattern(), DeclVisibilityKind::LocalVariable);`
`3115`	`3115`	`}`
`3116`	`3116`
Original file line number	Diff line number	Diff line change
`@@ -1092,7 +1092,7 @@ class ExprContextAnalyzer {`
`1092`	`1092`	`break;`
`1093`	`1093`	`}`
`1094`	`1094`	`case StmtKind::ForEach:`
`1095`		`- if (auto SEQ = cast<ForEachStmt>(Parent)->getSequence()) {`
	`1095`	`+ if (auto SEQ = cast<ForEachStmt>(Parent)->getParsedSequence()) {`
`1096`	`1096`	`if (containsTarget(SEQ)) {`
`1097`	`1097`	`recordPossibleType(Context.getSequenceType());`
`1098`	`1098`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2246,7 +2246,7 @@ class FormatWalker : public ASTWalker {`
`2246`	`2246`	`if (Range.isValid() && overlapsTarget(Range))`
`2247`	`2247`	`return IndentContext {ForLoc, !OutdentChecker::hasOutdent(SM, P)};`
`2248`	`2248`	`}`
`2249`		`- if (auto *E = FS->getSequence()) {`
	`2249`	`+ if (auto *E = FS->getParsedSequence()) {`
`2250`	`2250`	`SourceRange Range = FS->getInLoc();`
`2251`	`2251`	`widenOrSet(Range, E->getSourceRange());`
`2252`	`2252`	`if (Range.isValid() && isTargetContext(Range)) {`
Original file line number	Diff line number	Diff line change
`@@ -747,8 +747,8 @@ std::pair<bool, Stmt > ModelASTWalker::walkToStmtPre(Stmt S) {`
`747`	`747`	`charSourceRangeFromSourceRange(SM, ElemRange));`
`748`	`748`	`}`
`749`	`749`	`}`
`750`		`- if (ForEachS->getSequence())`
`751`		`- addExprElem(SyntaxStructureElementKind::Expr, ForEachS->getSequence(),SN);`
	`750`	`+ if (auto *S = ForEachS->getParsedSequence())`
	`751`	`+ addExprElem(SyntaxStructureElementKind::Expr, S, SN);`
`752`	`752`	`pushStructureNode(SN, S);`
`753`	`753`
`754`	`754`	`} else if (auto *WhileS = dyn_cast<WhileStmt>(S)) {`