[Function builders] Add code completion for build* functions.

When performing code completion inside the declaration of a type with
the function builder attribute, also include completions for all of
the possible "build" functions.

Author: DougGregor

include/swift/Sema/IDETypeChecking.h

@@ -254,6 +254,29 @@ namespace swift {
                             bool IncludeProtocolRequirements = true,
                             bool Transitive = false);
 
+  /// Enumerates the various kinds of "build" functions within a function
+  /// builder.
+  enum class FunctionBuilderBuildFunction {
+    BuildBlock,
+    BuildExpression,
+    BuildOptional,
+    BuildEitherFirst,
+    BuildEitherSecond,
+    BuildArray,
+    BuildLimitedAvailability,
+    BuildFinalResult,
+  };
+
+  /// Try to infer the component type of a function builder from the type
+  /// of buildBlock or buildExpression, if it was there.
+  Type inferFunctionBuilderComponentType(NominalTypeDecl *builder);
+
+  /// Print the declaration for a function builder "build" function, for use
+  /// in Fix-Its, code completion, and so on.
+  void printFunctionBuilderBuildFunction(
+      NominalTypeDecl *builder, Type componentType,
+      FunctionBuilderBuildFunction function,
+      Optional<std::string> stubIndent, llvm::raw_ostream &out);
 }
 
 #endif

lib/IDE/CodeCompletion.cpp

@@ -5136,6 +5136,60 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
     }
   }
 
+  void addFunctionBuilderBuildCompletion(
+      NominalTypeDecl *builder, Type componentType,
+      FunctionBuilderBuildFunction function) {
+    CodeCompletionResultBuilder Builder(
+        Sink,
+        CodeCompletionResult::ResultKind::Pattern,
+        SemanticContextKind::CurrentNominal, {});
+    Builder.setExpectedTypeRelation(
+        CodeCompletionResult::ExpectedTypeRelation::NotApplicable);
+
+    if (!hasFuncIntroducer) {
+      if (!hasAccessModifier &&
+          builder->getFormalAccess() >= AccessLevel::Public)
+        Builder.addAccessControlKeyword(AccessLevel::Public);
+
+      if (!hasStaticOrClass)
+        Builder.addKeyword("static");
+
+      Builder.addKeyword("func");
+    }
+
+    std::string declStringWithoutFunc;
+    {
+      llvm::raw_string_ostream out(declStringWithoutFunc);
+      printFunctionBuilderBuildFunction(
+          builder, componentType, function, None, out);
+    }
+    Builder.addTextChunk(declStringWithoutFunc);
+    Builder.addBraceStmtWithCursor();
+  }
+
+  /// Add completions for the various "build" functions in a function builder.
+  void addFunctionBuilderBuildCompletions(NominalTypeDecl *builder) {
+    Type componentType = inferFunctionBuilderComponentType(builder);
+
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildBlock);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildExpression);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildOptional);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildEitherFirst);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildEitherSecond);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildArray);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType,
+        FunctionBuilderBuildFunction::BuildLimitedAvailability);
+    addFunctionBuilderBuildCompletion(
+        builder, componentType, FunctionBuilderBuildFunction::BuildFinalResult);
+  }
+
   void getOverrideCompletions(SourceLoc Loc) {
     if (!CurrDeclContext->isTypeContext())
       return;
@@ -5154,6 +5208,10 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
       addDesignatedInitializers(NTD);
       addAssociatedTypes(NTD);
     }
+
+    if (NTD && NTD->getAttrs().hasAttribute<FunctionBuilderAttr>()) {
+      addFunctionBuilderBuildCompletions(NTD);
+    }
   }
 };
 } // end anonymous namespace

lib/Sema/BuilderTransform.cpp

@@ -20,6 +20,8 @@
 #include "SolutionResult.h"
 #include "TypeChecker.h"
 #include "TypeCheckAvailability.h"
+#include "swift/Sema/IDETypeChecking.h"
+#include "swift/AST/ASTPrinter.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/NameLookup.h"
@@ -1884,3 +1886,107 @@ bool TypeChecker::typeSupportsBuilderOp(
   return foundMatch;
 }
 
+Type swift::inferFunctionBuilderComponentType(NominalTypeDecl *builder) {
+  Type componentType;
+
+  SmallVector<ValueDecl *, 4> potentialMatches;
+  ASTContext &ctx = builder->getASTContext();
+  bool supportsBuildBlock = TypeChecker::typeSupportsBuilderOp(
+      builder->getDeclaredInterfaceType(), builder, ctx.Id_buildBlock,
+      /*argLabels=*/{}, &potentialMatches);
+  if (supportsBuildBlock) {
+    for (auto decl : potentialMatches) {
+      auto func = dyn_cast<FuncDecl>(decl);
+      if (!func || !func->isStatic())
+        continue;
+
+      // If we haven't seen a component type before, gather it.
+      if (!componentType) {
+        componentType = func->getResultInterfaceType();
+        continue;
+      }
+
+      // If there are inconsistent component types, bail out.
+      if (!componentType->isEqual(func->getResultInterfaceType())) {
+        componentType = Type();
+        break;
+      }
+    }
+  }
+
+  return componentType;
+}
+
+void swift::printFunctionBuilderBuildFunction(
+      NominalTypeDecl *builder, Type componentType,
+      FunctionBuilderBuildFunction function,
+      Optional<std::string> stubIndent, llvm::raw_ostream &out) {
+  // Render the component type into a string.
+  std::string componentTypeString;
+  if (componentType)
+    componentTypeString = componentType.getString();
+  else
+    componentTypeString = "<#Component#>";
+
+  // Render the code.
+  ExtraIndentStreamPrinter printer(out, stubIndent.getValueOr(std::string()));
+
+  // If we're supposed to provide a full stub, add a newline and the introducer
+  // keywords.
+  if (stubIndent) {
+    printer.printNewline();
+
+    if (builder->getFormalAccess() >= AccessLevel::Public)
+      printer << "public ";
+
+    printer << "static func ";
+  }
+
+  bool printedResult = false;
+  switch (function) {
+  case FunctionBuilderBuildFunction::BuildBlock:
+    printer << "buildBlock(_ components: " << componentTypeString << "...)";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildExpression:
+    printer << "buildExpression(_ expression: <#Expression#>)";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildOptional:
+    printer << "buildOptional(_ component: " << componentTypeString << "?)";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildEitherFirst:
+    printer << "buildEither(first component: " << componentTypeString << ")";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildEitherSecond:
+    printer << "buildEither(second component: " << componentTypeString << ")";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildArray:
+    printer << "buildArray(_ components: [" << componentTypeString << "])";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildLimitedAvailability:
+    printer << "buildLimitedAvailability(_ component: " << componentTypeString
+            << ")";
+    break;
+
+  case FunctionBuilderBuildFunction::BuildFinalResult:
+    printer << "buildFinalResult(_ component: " << componentTypeString
+            << ") -> <#Result#>";
+    printedResult = true;
+    break;
+  }
+
+  if (!printedResult)
+    printer << " -> " << componentTypeString << " {";
+
+  if (stubIndent) {
+    printer.printNewline();
+    printer << "  <#code#>";
+    printer.printNewline();
+    printer << "}";
+  }
+}

lib/Sema/CSDiagnostics.cpp

@@ -19,6 +19,7 @@
 #include "MiscDiagnostics.h"
 #include "TypeCheckProtocol.h"
 #include "TypoCorrection.h"
+#include "swift/Sema/IDETypeChecking.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTPrinter.h"
 #include "swift/AST/Decl.h"
@@ -5501,96 +5502,6 @@ static bool hasMissingElseInChain(IfStmt *ifStmt) {
   return false;
 }
 
-namespace  {
-  enum class MissingBuildFunction {
-    BuildOptional,
-    BuildEitherFirst,
-    BuildEitherSecond,
-    BuildArray,
-  };
-}
-
-/// Add code for a missing 'build' function to a function builder.
-static void printMissingBuildFunctionCode(
-    SourceLoc buildInsertionLoc, NominalTypeDecl *builder,
-    MissingBuildFunction missing, llvm::raw_ostream &out) {
-  // Determine the component type from the builder itself, if we can.
-  ASTContext &ctx = builder->getASTContext();
-  Type componentType;
-  bool isPublic = false;
-  {
-    SmallVector<ValueDecl *, 4> potentialMatches;
-    bool supportsBuildBlock = TypeChecker::typeSupportsBuilderOp(
-        builder->getDeclaredInterfaceType(), builder, ctx.Id_buildBlock,
-        /*argLabels=*/{}, &potentialMatches);
-    if (supportsBuildBlock) {
-      for (auto decl : potentialMatches) {
-        auto func = dyn_cast<FuncDecl>(decl);
-        if (!func || !func->isStatic())
-          continue;
-
-        // If we haven't seen a component type before, gather it.
-        if (!componentType) {
-          componentType = func->getResultInterfaceType();
-          isPublic = func->getFormalAccess() == AccessLevel::Public;
-          continue;
-        }
-
-        // If there are inconsistent component types, bail out.
-        if (!componentType->isEqual(func->getResultInterfaceType())) {
-          componentType = Type();
-          isPublic = false;
-          break;
-        }
-      }
-    }
-  }
-
-  // Render the component type into a string.
-  std::string componentTypeString;
-  if (componentType)
-    componentTypeString = componentType.getString();
-  else
-    componentTypeString = "<#Component#>";
-
-  // Render the code.
-  StringRef extraIndent;
-  StringRef currentIndent = Lexer::getIndentationForLine(
-      ctx.SourceMgr, buildInsertionLoc, &extraIndent);
-  std::string stubIndent = (currentIndent + extraIndent).str();
-
-  ExtraIndentStreamPrinter printer(out, stubIndent);
-  printer.printNewline();
-
-  if (isPublic)
-    printer << "public ";
-
-  printer << "static func ";
-  switch (missing) {
-  case MissingBuildFunction::BuildOptional:
-    printer << "buildOptional(_ component: " << componentTypeString << "?)";
-    break;
-
-  case MissingBuildFunction::BuildEitherFirst:
-    printer << "buildEither(first component: " << componentTypeString << ")";
-    break;
-
-  case MissingBuildFunction::BuildEitherSecond:
-    printer << "buildEither(second component: " << componentTypeString << ")";
-    break;
-
-  case MissingBuildFunction::BuildArray:
-    printer << "buildArray(_ components: [" << componentTypeString << "])";
-    break;
-  }
-
-  printer << " -> " << componentTypeString << " {";
-  printer.printNewline();
-  printer << "  <#code#>";
-  printer.printNewline();
-  printer << "}";
-}
-
 void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
     bool asNote) {
   if (auto stmt = unhandled.dyn_cast<Stmt *>()) {
@@ -5601,9 +5512,19 @@ void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
     // Emit custom notes to help the user introduce the appropriate 'build'
     // functions.
     SourceLoc buildInsertionLoc = builder->getBraces().Start;
+    std::string stubIndent;
+    Type componentType;
     if (buildInsertionLoc.isValid()) {
       buildInsertionLoc = Lexer::getLocForEndOfToken(
           getASTContext().SourceMgr, buildInsertionLoc);
+
+      ASTContext &ctx = getASTContext();
+      StringRef extraIndent;
+      StringRef currentIndent = Lexer::getIndentationForLine(
+          ctx.SourceMgr, buildInsertionLoc, &extraIndent);
+      stubIndent = (currentIndent + extraIndent).str();
+
+      componentType = inferFunctionBuilderComponentType(builder);
     }
 
     if (buildInsertionLoc.isInvalid()) {
@@ -5616,9 +5537,9 @@ void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
       std::string fixItString;
       {
         llvm::raw_string_ostream out(fixItString);
-        printMissingBuildFunctionCode(
-            buildInsertionLoc, builder, MissingBuildFunction::BuildOptional,
-            out);
+        printFunctionBuilderBuildFunction(
+            builder, componentType, FunctionBuilderBuildFunction::BuildOptional,
+            stubIndent, out);
       }
 
       diag.fixItInsert(buildInsertionLoc, fixItString);
@@ -5630,13 +5551,15 @@ void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
       std::string fixItString;
       {
         llvm::raw_string_ostream out(fixItString);
-        printMissingBuildFunctionCode(
-            buildInsertionLoc, builder, MissingBuildFunction::BuildEitherFirst,
-            out);
+        printFunctionBuilderBuildFunction(
+            builder, componentType,
+            FunctionBuilderBuildFunction::BuildEitherFirst,
+            stubIndent, out);
         out << '\n';
-        printMissingBuildFunctionCode(
-            buildInsertionLoc, builder, MissingBuildFunction::BuildEitherSecond,
-            out);
+        printFunctionBuilderBuildFunction(
+            builder, componentType,
+            FunctionBuilderBuildFunction::BuildEitherSecond,
+            stubIndent, out);
       }
 
       diag.fixItInsert(buildInsertionLoc, fixItString);
@@ -5648,8 +5571,9 @@ void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
       std::string fixItString;
       {
         llvm::raw_string_ostream out(fixItString);
-        printMissingBuildFunctionCode(
-            buildInsertionLoc, builder, MissingBuildFunction::BuildArray, out);
+        printFunctionBuilderBuildFunction(
+            builder, componentType, FunctionBuilderBuildFunction::BuildArray,
+            stubIndent, out);
       }
 
       diag.fixItInsert(buildInsertionLoc, fixItString);