Merge pull request #37760 from hamishknight/walk-this-way

[ASTWalker] Walk ForEachStmt in source order

Author: hamishknight

lib/AST/ASTWalker.cpp

@@ -1563,13 +1563,6 @@ Stmt *Traversal::visitForEachStmt(ForEachStmt *S) {
       return nullptr;
   }
 
-  if (Expr *Where = S->getWhere()) {
-    if ((Where = doIt(Where)))
-      S->setWhere(Where);
-    else
-      return nullptr;
-  }
-
   // The iterator decl is built directly on top of the sequence
   // expression, so don't visit both.
   if (Expr *Sequence = S->getSequence()) {
@@ -1579,6 +1572,13 @@ Stmt *Traversal::visitForEachStmt(ForEachStmt *S) {
       return nullptr;
   }
 
+  if (Expr *Where = S->getWhere()) {
+    if ((Where = doIt(Where)))
+      S->setWhere(Where);
+    else
+      return nullptr;
+  }
+
   if (auto IteratorNext = S->getConvertElementExpr()) {
     if ((IteratorNext = doIt(IteratorNext)))
       S->setConvertElementExpr(IteratorNext);

test/refactoring/ConvertAsync/basic.swift

@@ -855,4 +855,22 @@ class TestConvertFunctionWithCallToFunctionsWithSpecialName {
     _ = x[1]
     return x
   }
-}
+}
+
+// rdar://78781061
+// RUN: %refactor -convert-to-async -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=FOR-IN-WHERE %s
+func testForInWhereRefactoring() {
+  let arr: [String] = []
+  for str in arr where str.count != 0 {
+    simple { res in
+      print(res)
+    }
+  }
+}
+// FOR-IN-WHERE:      func testForInWhereRefactoring() async {
+// FOR-IN-WHERE-NEXT:   let arr: [String] = []
+// FOR-IN-WHERE-NEXT:   for str in arr where str.count != 0 {
+// FOR-IN-WHERE-NEXT:     let res = await simple()
+// FOR-IN-WHERE-NEXT:     print(res)
+// FOR-IN-WHERE-NEXT:   }
+// FOR-IN-WHERE-NEXT: }