wip: desugar ForEachStmt at AST level

I have not yet tested any behavior of generated SIL code but it seems
reasonable in comparison.

I was able to compile the standard library.

Here:
* include/swift/AST/ASTBridging.h,
* include/swift/AST/Expr.h,
* include/swift/AST/ExprNodes.def,
* include/swift/AST/Stmt.h,
* include/swift/AST/StmtNodes.def,
* include/swift/AST/TypeCheckRequests.h,
* include/swift/AST/TypeCheckerTypeIDZone.def,
* include/swift/Sema/ConstraintLocator.h,
* include/swift/Sema/SyntacticElementTarget.h,
* lib/AST/ASTDumper.cpp,
* lib/AST/ASTPrinter.cpp,
* lib/AST/ASTScopeCreation.cpp,
* lib/AST/ASTVerifier.cpp,
* lib/AST/ASTWalker.cpp,
* lib/AST/Bridging/StmtBridging.cpp,
* lib/AST/Expr.cpp,
* lib/AST/Stmt.cpp,
* lib/AST/TypeCheckRequests.cpp,
* lib/ASTGen/Sources/ASTGen/Stmts.swift,
* lib/Parse/ParseStmt.cpp,
* lib/SILGen/ASTVisitor.h,
* lib/SILGen/SILGenExpr.cpp,
* lib/SILGen/SILGenStmt.cpp,
* lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.cpp,
* lib/Sema/BuilderTransform.cpp,
* lib/Sema/CSApply.cpp,
* lib/Sema/CSDiagnostics.cpp,
* lib/Sema/CSGen.cpp,
* lib/Sema/CSSimplify.cpp,
* lib/Sema/CSSyntacticElement.cpp,
* lib/Sema/SyntacticElementTarget.cpp,
* lib/Sema/TypeCheckEffects.cpp,
* lib/Sema/TypeCheckStmt.cpp.

Author: elsakeirouz

include/swift/AST/ASTBridging.h

@@ -2404,13 +2404,14 @@ BridgedFallthroughStmt_createParsed(swift::SourceLoc loc,
                                     BridgedDeclContext cDC);
 
 SWIFT_NAME("BridgedForEachStmt.createParsed(_:labelInfo:forLoc:tryLoc:awaitLoc:"
-           "unsafeLoc:pattern:inLoc:sequence:whereLoc:whereExpr:body:)")
+           "unsafeLoc:pattern:inLoc:sequence:whereLoc:whereExpr:body:declContext:)")
 BridgedForEachStmt BridgedForEachStmt_createParsed(
     BridgedASTContext cContext, BridgedLabeledStmtInfo cLabelInfo,
     swift::SourceLoc forLoc, swift::SourceLoc tryLoc, swift::SourceLoc awaitLoc,
     swift::SourceLoc unsafeLoc, BridgedPattern cPat, swift::SourceLoc inLoc,
     BridgedExpr cSequence, swift::SourceLoc whereLoc,
-    BridgedNullableExpr cWhereExpr, BridgedBraceStmt cBody);
+    BridgedNullableExpr cWhereExpr, BridgedBraceStmt cBody,
+    BridgedDeclContext cDeclContext);
 
 SWIFT_NAME("BridgedGuardStmt.createParsed(_:guardLoc:conds:body:)")
 BridgedGuardStmt BridgedGuardStmt_createParsed(BridgedASTContext cContext,

include/swift/AST/Expr.h

@@ -6724,6 +6724,27 @@ class MacroExpansionExpr final : public Expr,
   }
 };
 
+/// OpaqueExpr - created to serve as an indirection to a ForEachStmt's sequence
+/// expr and where clause to avoid visiting it twice in the ASTWalker after
+/// having desugared the loop. This will only be processed in SILGen to emit
+/// the underlying expression.
+class OpaqueExpr final : public Expr {
+  Expr *OriginalExpr;
+
+public:
+  OpaqueExpr(Expr* originalExpr)
+    : Expr(ExprKind::Opaque, /*implicit*/ true, originalExpr->getType()),
+      OriginalExpr(originalExpr) {}
+
+  Expr *getOriginalExpr() const { return OriginalExpr; }
+  SourceLoc getStartLoc() const { return OriginalExpr->getStartLoc(); }
+  SourceLoc getEndLoc() const { return OriginalExpr->getEndLoc(); }
+
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::Opaque;
+  }
+};
+
 inline bool Expr::isInfixOperator() const {
   return isa<BinaryExpr>(this) || isa<TernaryExpr>(this) ||
          isa<AssignExpr>(this) || isa<ExplicitCastExpr>(this);

include/swift/AST/ExprNodes.def

@@ -218,6 +218,7 @@ EXPR(Tap, Expr)
 UNCHECKED_EXPR(TypeJoin, Expr)
 EXPR(MacroExpansion, Expr)
 EXPR(TypeValue, Expr)
+EXPR(Opaque, Expr)
 // Don't forget to update the LAST_EXPR below when adding a new Expr here.
 LAST_EXPR(TypeValue)
 

include/swift/AST/Stmt.h

@@ -1003,21 +1003,22 @@ class ForEachStmt : public LabeledStmt {
   SourceLoc WhereLoc;
   Expr *WhereExpr = nullptr;
   BraceStmt *Body;
+  DeclContext* DC = nullptr;
 
   // Set by Sema:
   ProtocolConformanceRef sequenceConformance = ProtocolConformanceRef();
   Type sequenceType;
-  PatternBindingDecl *iteratorVar = nullptr;
-  Expr *nextCall = nullptr;
-  OpaqueValueExpr *elementExpr = nullptr;
+  BraceStmt *desugaredStmt = nullptr;
   Expr *convertElementExpr = nullptr;
+  LabeledStmt *continueTarget = nullptr;
+  LabeledStmt *breakTarget = nullptr;
 
 public:
   ForEachStmt(LabeledStmtInfo LabelInfo, SourceLoc ForLoc, SourceLoc TryLoc,
               SourceLoc AwaitLoc, SourceLoc UnsafeLoc, Pattern *Pat,
               SourceLoc InLoc, Expr *Sequence,
               SourceLoc WhereLoc, Expr *WhereExpr, BraceStmt *Body,
-              std::optional<bool> implicit = std::nullopt)
+              DeclContext* DC, std::optional<bool> implicit = std::nullopt)
       : LabeledStmt(StmtKind::ForEach, getDefaultImplicitFlag(implicit, ForLoc),
                     LabelInfo),
         ForLoc(ForLoc), TryLoc(TryLoc), AwaitLoc(AwaitLoc), UnsafeLoc(UnsafeLoc),
@@ -1026,15 +1027,6 @@ class ForEachStmt : public LabeledStmt {
     setPattern(Pat);
   }
 
-  void setIteratorVar(PatternBindingDecl *var) { iteratorVar = var; }
-  PatternBindingDecl *getIteratorVar() const { return iteratorVar; }
-
-  void setNextCall(Expr *next) { nextCall = next; }
-  Expr *getNextCall() const { return nextCall; }
-
-  void setElementExpr(OpaqueValueExpr *expr) { elementExpr = expr; }
-  OpaqueValueExpr *getElementExpr() const { return elementExpr; }
-
   void setConvertElementExpr(Expr *expr) { convertElementExpr = expr; }
   Expr *getConvertElementExpr() const { return convertElementExpr; }
 
@@ -1076,20 +1068,29 @@ class ForEachStmt : public LabeledStmt {
   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; }
   void setBody(BraceStmt *B) { Body = B; }
 
   SourceLoc getStartLoc() const { return getLabelLocOrKeywordLoc(ForLoc); }
   SourceLoc getEndLoc() const { return Body->getEndLoc(); }
+
+  DeclContext *getDeclContext() const { return DC; }
+  void setDeclContext(DeclContext *newDC) { DC = newDC; }
 
   static bool classof(const Stmt *S) {
     return S->getKind() == StmtKind::ForEach;
   }
+
+  BraceStmt* desugar();
+  BraceStmt* getDesugaredStmt() const { return desugaredStmt; }
+  void setDesugaredStmt(BraceStmt* newStmt) { desugaredStmt = newStmt; }
+
+  void setContinueTarget(LabeledStmt *target) { continueTarget = target; }
+  LabeledStmt* getContinueTarget() { return continueTarget; }
+
+  void setBreakTarget(LabeledStmt *target) { breakTarget = target; }
+  LabeledStmt* getBreakTarget() { return breakTarget; }
 };
 
 /// A pattern and an optional guard expression used in a 'case' statement.
@@ -1541,6 +1542,31 @@ class DoCatchStmt final
   }
 };
 
+/// OpaqueStmt - created to serve as an indirection to a ForEachStmt's body
+/// to avoid visiting it twice in the ASTWalker after having desugared the loop.
+/// This ensures we only visit the body once, and this OpaqueStmt will only be
+/// visited to emit the underlying statement in SILGen.
+class OpaqueStmt final : public Stmt {
+  SourceLoc StartLoc;
+  SourceLoc EndLoc;
+  BraceStmt *Body; // FIXME: should I just use Stmt * so that this is more versatile?
+                   // If not, should the class be renamed to be more specific?
+  public:
+    OpaqueStmt(BraceStmt* body, SourceLoc startLoc, SourceLoc endLoc)
+    : Stmt(StmtKind::Opaque, true /*always implicit*/),
+      StartLoc(startLoc), EndLoc(endLoc), Body(body) {}
+
+  SourceLoc getLoc() const { return StartLoc; }
+  SourceLoc getStartLoc() const { return StartLoc; }
+  SourceLoc getEndLoc() const { return EndLoc; }
+
+  BraceStmt* getUnderlyingStmt() { return Body; }
+
+  static bool classof(const Stmt *S) {
+    return S->getKind() == StmtKind::Opaque;
+  }
+};
+
 /// BreakStmt - The "break" and "break label" statement.
 class BreakStmt : public Stmt {
   SourceLoc Loc;

include/swift/AST/StmtNodes.def

@@ -61,6 +61,7 @@ ABSTRACT_STMT(Labeled, Stmt)
   LABELED_STMT(ForEach, LabeledStmt)
   LABELED_STMT(Switch, LabeledStmt)
   STMT_RANGE(Labeled, If, Switch)
+STMT(Opaque, Stmt)
 STMT(Case, Stmt)
 STMT(Break, Stmt)
 STMT(Continue, Stmt)

include/swift/AST/TypeCheckRequests.h

@@ -5591,6 +5591,25 @@ class IsCustomAvailabilityDomainPermanentlyEnabled
   }
 };
 
+class DesugarForEachStmtRequest
+    : public SimpleRequest<DesugarForEachStmtRequest,
+                           BraceStmt *(ForEachStmt*),
+                           RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  BraceStmt *evaluate(Evaluator &evaluator, ForEachStmt *FES) const;
+
+public:
+  bool isCached() const { return true; }
+  std::optional<BraceStmt*> getCachedResult() const;
+  void cacheResult(BraceStmt *stmt) const;
+};
+
 #define SWIFT_TYPEID_ZONE TypeChecker
 #define SWIFT_TYPEID_HEADER "swift/AST/TypeCheckerTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -674,3 +674,7 @@ SWIFT_REQUEST(TypeChecker, IsCustomAvailabilityDomainPermanentlyEnabled,
 SWIFT_REQUEST(TypeChecker, EmitPerformanceHints,
               evaluator::SideEffect(SourceFile *),
               Cached, NoLocationInfo)
+
+SWIFT_REQUEST(TypeChecker, DesugarForEachStmtRequest,
+              Stmt*(const ForEachStmt*),
+              Cached, NoLocationInfo)

include/swift/Sema/ConstraintLocator.h

@@ -83,6 +83,8 @@ enum ContextualTypePurpose : uint8_t {
 
   CTP_ExprPattern,      ///< `~=` operator application associated with expression
                         /// pattern.
+
+  CTP_ForEachElement,    ///< Element expression associated with `for-in` loop.
 };
 
 namespace constraints {

include/swift/Sema/SyntacticElementTarget.h

@@ -41,12 +41,6 @@ struct SequenceIterationInfo {
 
   /// The type of the pattern that matches the elements.
   Type initType;
-
-  /// Implicit `$iterator = <sequence>.makeIterator()`
-  PatternBindingDecl *makeIteratorVar;
-
-  /// Implicit `$iterator.next()` call.
-  Expr *nextCall;
 };
 
 /// Describes information about a for-in loop over a pack that needs to be
@@ -605,6 +599,7 @@ class SyntacticElementTarget {
     case CTP_Initialization:
     case CTP_ForEachSequence:
     case CTP_ExprPattern:
+    case CTP_ForEachElement:
       break;
     default:
       assert(false && "Unexpected contextual type purpose");

lib/AST/ASTDumper.cpp

@@ -3252,6 +3252,10 @@ class PrintStmt : public StmtVisitor<PrintStmt, void, Label>,
     printFlag(S->TrailingSemiLoc.isValid(), "trailing_semi");
   }
 
+  void visitOpaqueStmt(OpaqueStmt *S, Label label){
+    visitBraceStmt(S->getUnderlyingStmt(), label);
+  }
+
   void visitBraceStmt(BraceStmt *S, Label label) {
     printCommon(S, "brace_stmt", label);
     printList(S->getElements(), [&](auto &Elt, Label label) {
@@ -3332,20 +3336,15 @@ class PrintStmt : public StmtVisitor<PrintStmt, void, Label>,
       printRec(S->getWhere(), Label::always("where"));
     }
     printRec(S->getParsedSequence(), Label::optional("parsed_sequence"));
-    if (S->getIteratorVar()) {
-      printRec(S->getIteratorVar(), Label::optional("iterator_var"));
-    }
-    if (S->getNextCall()) {
-      printRec(S->getNextCall(), Label::optional("next_call"));
-    }
     if (S->getConvertElementExpr()) {
       printRec(S->getConvertElementExpr(),
                Label::optional("convert_element_expr"));
     }
-    if (S->getElementExpr()) {
-      printRec(S->getElementExpr(), Label::optional("element_expr"));
-    }
+
     printRec(S->getBody(), Label::optional("body"));
+
+    printRec(S->getDesugaredStmt(), Label::optional("desugared_loop"));
+
     printFoot();
   }
   void visitBreakStmt(BreakStmt *S, Label label) {
@@ -4237,6 +4236,10 @@ class PrintExpr : public ExprVisitor<PrintExpr, void, Label>,
     printFoot();
   }
 
+  void visitOpaqueExpr(OpaqueExpr *E, Label label){
+    visit(E->getOriginalExpr(), label);
+  }
+
   void visitPropertyWrapperValuePlaceholderExpr(
       PropertyWrapperValuePlaceholderExpr *E, Label label) {
     printCommon(E, "property_wrapper_value_placeholder_expr", label);