[5.5] Visit ForEachStmt in source order in SemaAnnotator

lib/IDE/SourceEntityWalker.cpp

@@ -234,6 +234,65 @@ std::pair<bool, Stmt *> SemaAnnotator::walkToStmtPre(Stmt *S) {
       // Already walked children.
       return { false, Continue ? DeferS : nullptr };
     }
+
+    auto doSkipChildren = [&]() -> std::pair<bool, Stmt *> {
+      if (!walkToStmtPost(S))
+        return {false, nullptr};
+      return {false, S};
+    };
+
+    auto doStopTraversal = [&]() -> std::pair<bool, Stmt *> {
+      Cancelled = true;
+      return {false, nullptr};
+    };
+
+    // Make sure to walk a ForEachStmt in source order. Note this is a narrow
+    // fix for release/5.5. On main, this is fixed in the ASTWalker itself.
+    // rdar://78781061
+    if (auto *FE = dyn_cast<ForEachStmt>(S)) {
+      if (auto *P = FE->getPattern()) {
+        auto *NewP = P->walk(*this);
+        if (!NewP)
+          return doStopTraversal();
+        assert(NewP == P);
+      }
+      if (auto *SE = FE->getSequence()) {
+        auto *NewSE = SE->walk(*this);
+        if (!NewSE)
+          return doStopTraversal();
+        assert(NewSE == SE);
+      }
+      if (auto *Where = FE->getWhere()) {
+        auto *NewWhere = Where->walk(*this);
+        if (!NewWhere)
+          return doStopTraversal();
+        assert(NewWhere == Where);
+      }
+      if (auto *IteratorNext = FE->getConvertElementExpr()) {
+        auto *NewIteratorNext = IteratorNext->walk(*this);
+        if (!NewIteratorNext)
+          return doStopTraversal();
+        assert(NewIteratorNext == IteratorNext);
+      }
+      if (auto *IteratorVar = FE->getIteratorVar()) {
+        if (IteratorVar->walk(*this))
+          return doStopTraversal();
+      }
+      if (auto *IteratorVarRef = FE->getIteratorVarRef()) {
+        auto *NewIteratorVarRef = IteratorVarRef->walk(*this);
+        if (!NewIteratorVarRef)
+          return doStopTraversal();
+        assert(NewIteratorVarRef == IteratorVarRef);
+      }
+      if (auto *Body = FE->getBody()) {
+        auto *NewBody = Body->walk(*this);
+        if (!NewBody)
+          return doStopTraversal();
+        assert(NewBody == Body);
+      }
+      // Already visited.
+      return doSkipChildren();
+    }
   }
   return { TraverseChildren, S };
 }

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