[PreCheckTarget] Move closure body expansion into PreCheckTarget::walkToClosureExprPre

If closure body is expanded in expression context we need to do
some special handling to make sure that source ranges are not
completely broken.

lib/Sema/PreCheckTarget.cpp

@@ -1149,7 +1149,7 @@ public:
 
   ASTContext &getASTContext() const { return Ctx; }
 
-  bool walkToClosureExprPre(ClosureExpr *expr);
+  bool walkToClosureExprPre(ClosureExpr *expr, ParentTy &parent);
 
   MacroWalking getMacroWalkingBehavior() const override {
     return MacroWalking::Arguments;
@@ -1251,7 +1251,7 @@ public:
     // but do not walk into the body. That will be type-checked after
     // we've determine the complete function type.
     if (auto closure = dyn_cast<ClosureExpr>(expr))
-      return finish(walkToClosureExprPre(closure), expr);
+      return finish(walkToClosureExprPre(closure, Parent), expr);
 
     if (auto *unresolved = dyn_cast<UnresolvedDeclRefExpr>(expr))
       return finish(true, TypeChecker::resolveDeclRefExpr(unresolved, DC));
@@ -1533,7 +1533,27 @@ public:
 
 /// Perform prechecking of a ClosureExpr before we dive into it.  This returns
 /// true when we want the body to be considered part of this larger expression.
-bool PreCheckTarget::walkToClosureExprPre(ClosureExpr *closure) {
+bool PreCheckTarget::walkToClosureExprPre(ClosureExpr *closure,
+                                          ParentTy &parent) {
+  if (auto *expandedBody = closure->getExpandedBody()) {
+    if (Parent.getAsExpr()) {
+      // We cannot simply replace the body when closure i.e. is passed
+      // as an argument to a call or is a source of an assignment
+      // because the source range of the argument list would cross
+      // buffer boundaries. One way to avoid that is to inject
+      // elements into a new implicit brace statement with the original
+      // source locations. Brace statement has to be implicit because its
+      // elements are in a different buffer.
+      auto sourceRange = closure->getSourceRange();
+      closure->setBody(BraceStmt::create(getASTContext(), sourceRange.Start,
+                                         expandedBody->getElements(),
+                                         sourceRange.End,
+                                         /*implicit=*/true));
+    } else {
+      closure->setBody(expandedBody);
+    }
+  }
+
   // Pre-check the closure body.
   (void)evaluateOrDefault(Ctx.evaluator, PreCheckClosureBodyRequest{closure},
                           nullptr);