[CodeCompletion] Fast completion inside function builder function

include/swift/AST/Stmt.h

@@ -1353,6 +1353,13 @@ class PoundAssertStmt : public Stmt {
   }
 };
 
+inline void simple_display(llvm::raw_ostream &out, Stmt *S) {
+  if (S)
+    out << Stmt::getKindName(S->getKind());
+  else
+    out << "(null)";
+};
+
 } // end namespace swift
 
 #endif // SWIFT_AST_STMT_H

include/swift/AST/TypeCheckRequests.h

@@ -1836,7 +1836,8 @@ enum class FunctionBuilderBodyPreCheck : uint8_t {
 
 class PreCheckFunctionBuilderRequest
     : public SimpleRequest<PreCheckFunctionBuilderRequest,
-                           FunctionBuilderBodyPreCheck(AnyFunctionRef),
+                           FunctionBuilderBodyPreCheck(AnyFunctionRef,
+                                                       BraceStmt *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -1845,8 +1846,8 @@ class PreCheckFunctionBuilderRequest
   friend SimpleRequest;
 
   // Evaluation.
-  FunctionBuilderBodyPreCheck
-  evaluate(Evaluator &evaluator, AnyFunctionRef fn) const;
+  FunctionBuilderBodyPreCheck evaluate(Evaluator &evaluator, AnyFunctionRef fn,
+                                       BraceStmt *body) const;
 
 public:
   // Separate caching.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -231,7 +231,7 @@ SWIFT_REQUEST(TypeChecker, HasUserDefinedDesignatedInitRequest,
 SWIFT_REQUEST(TypeChecker, HasMemberwiseInitRequest,
               bool(StructDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, PreCheckFunctionBuilderRequest,
-              FunctionBuilderClosurePreCheck(AnyFunctionRef),
+              FunctionBuilderClosurePreCheck(AnyFunctionRef, BraceStmt *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ResolveImplicitMemberRequest,
               evaluator::SideEffect(NominalTypeDecl *, ImplicitMemberAction),

lib/Sema/BuilderTransform.cpp

@@ -1453,7 +1453,7 @@ Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
   // If we encountered an error or there was an explicit result type,
   // bail out and report that to the caller.
   auto &ctx = func->getASTContext();
-  auto request = PreCheckFunctionBuilderRequest{func};
+  auto request = PreCheckFunctionBuilderRequest{func, func->getBody()};
   switch (evaluateOrDefault(
               ctx.evaluator, request, FunctionBuilderBodyPreCheck::Error)) {
   case FunctionBuilderBodyPreCheck::Okay:
@@ -1581,7 +1581,7 @@ ConstraintSystem::matchFunctionBuilder(
 
   // Pre-check the body: pre-check any expressions in it and look
   // for return statements.
-  auto request = PreCheckFunctionBuilderRequest{fn};
+  auto request = PreCheckFunctionBuilderRequest{fn, fn.getBody()};
   switch (evaluateOrDefault(getASTContext().evaluator, request,
                             FunctionBuilderBodyPreCheck::Error)) {
   case FunctionBuilderBodyPreCheck::Okay:
@@ -1744,8 +1744,14 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
 }
 
 FunctionBuilderBodyPreCheck
-PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval,
-                                         AnyFunctionRef fn) const {
+PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval, AnyFunctionRef fn,
+                                         BraceStmt *body) const {
+  // NOTE: 'body' is passed only for the request evaluater caching key.
+  // Since source tooling (e.g. code completion) might replace the body,
+  // the function alone is not sufficient for the key.
+  assert(fn.getBody() == body &&
+         "body must be the current body of the function");
+
   // We don't want to do the precheck if it will already have happened in
   // the enclosing expression.
   bool skipPrecheck = false;

test/SourceKit/CodeComplete/complete_sequence_builderfunc.swift

@@ -0,0 +1,51 @@
+protocol Entity {}
+struct Empty: Entity {
+    var value: Void = ()
+}
+
+@_functionBuilder
+struct Builder {
+  static func buildBlock() ->  { Empty() }
+  static func buildBlock<T: Entity>(_ t: T) -> T { t }
+}
+
+struct MyValue {
+    var id: Int
+    var title: String
+}
+
+func build(_ arg: (MyValue) -> String) -> Empty { Empty() }
+
+struct MyStruct {
+  @Builder var body: some Entity {
+    build { value in
+      value./*HERE * 2*/
+    } /*HERE*/
+  }
+}
+
+// RUN: %sourcekitd-test \
+// RUN:   -req=complete -pos=22:13 %s -- %s == \
+// RUN:   -req=complete -pos=22:13 %s -- %s == \
+// RUN:   -req=complete -pos=23:7 %s -- %s \
+// RUN: | tee %t.result | %FileCheck %s
+
+// CHECK-LABEL: key.results: [
+// CHECK-DAG: key.description: "id"
+// CHECK-DAG: key.description: "title"
+// CHECK-DAG: key.description: "self"
+// CHECK: ]
+// CHECK-NOT: key.reusingastcontext: 1 
+
+// CHECK-LABEL: key.results: [
+// CHECK-DAG: key.description: "id"
+// CHECK-DAG: key.description: "title"
+// CHECK-DAG: key.description: "self"
+// CHECK: ]
+// CHECK: key.reusingastcontext: 1 
+
+// CHECK-LABEL: key.results: [
+// CHECK-DAG: key.description: "value"
+// CHECK: ]
+// CHECK: key.reusingastcontext: 1 
+