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

Author: xedin

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; }

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();

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());

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:

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));

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);
 }
 

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,

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());
         }

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)) {

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)) {