[Function builders] Fix-Its and code completion improvements

include/swift/AST/DiagnosticsSema.def

@@ -5312,6 +5312,19 @@ NOTE(function_builder_buildblock_enum_case, none,
 NOTE(function_builder_buildblock_not_static_method, none,
      "potential match 'buildBlock' is not a static method", ())
 
+NOTE(function_builder_missing_build_optional, none,
+     "add 'buildOptional(_:)' to the function builder %0 to add support for "
+     "'if' statements without an 'else'", (Type))
+NOTE(function_builder_missing_build_either, none,
+     "add 'buildEither(first:)' and 'buildEither(second:)' to the function "
+     "builder %0 to add support for 'if'-'else' and 'switch'", (Type))
+NOTE(function_builder_missing_build_array, none,
+     "add 'buildArray(_:)' to the function builder %0 to add support for "
+     "'for'..'in' loops", (Type))
+NOTE(function_builder_missing_build_limited_availability, none,
+     "add 'buildLimitedAvailability(_:)' to the function "
+     "builder %0 to erase type information for less-available types", (Type))
+
 //------------------------------------------------------------------------------
 // MARK: Tuple Shuffle Diagnostics
 //------------------------------------------------------------------------------

include/swift/AST/EducationalNotes.def

@@ -81,4 +81,15 @@ EDUCATIONAL_NOTES(type_cannot_conform, "protocol-type-non-conformance.md")
 EDUCATIONAL_NOTES(unlabeled_trailing_closure_deprecated,
                   "trailing-closure-matching.md")
 
+EDUCATIONAL_NOTES(function_builder_static_buildblock,
+                  "function-builder-methods.md")
+EDUCATIONAL_NOTES(function_builder_missing_limited_availability,
+                  "function-builder-methods.md")
+EDUCATIONAL_NOTES(function_builder_missing_build_optional,
+                  "function-builder-methods.md")
+EDUCATIONAL_NOTES(function_builder_missing_build_either,
+                  "function-builder-methods.md")
+EDUCATIONAL_NOTES(function_builder_missing_build_array,
+                  "function-builder-methods.md")
+
 #undef EDUCATIONAL_NOTES

include/swift/IDE/CodeCompletion.h

@@ -640,11 +640,13 @@ class CodeCompletionResult {
                        CodeCompletionString *CompletionString,
                        ExpectedTypeRelation TypeDistance,
                        CodeCompletionOperatorKind KnownOperatorKind =
-                           CodeCompletionOperatorKind::None)
+                           CodeCompletionOperatorKind::None,
+                       StringRef BriefDocComment = StringRef())
       : Kind(Kind), KnownOperatorKind(unsigned(KnownOperatorKind)),
         SemanticContext(unsigned(SemanticContext)), NotRecommended(false),
         NotRecReason(NotRecommendedReason::NoReason),
         NumBytesToErase(NumBytesToErase), CompletionString(CompletionString),
+        BriefDocComment(BriefDocComment),
         TypeDistance(TypeDistance) {
     assert(Kind != Declaration && "use the other constructor");
     assert(CompletionString);
@@ -664,12 +666,13 @@ class CodeCompletionResult {
                        SemanticContextKind SemanticContext,
                        unsigned NumBytesToErase,
                        CodeCompletionString *CompletionString,
-                       ExpectedTypeRelation TypeDistance)
+                       ExpectedTypeRelation TypeDistance,
+                       StringRef BriefDocComment = StringRef())
       : Kind(Keyword), KnownOperatorKind(0),
         SemanticContext(unsigned(SemanticContext)), NotRecommended(false),
         NotRecReason(NotRecommendedReason::NoReason),
         NumBytesToErase(NumBytesToErase), CompletionString(CompletionString),
-        TypeDistance(TypeDistance) {
+        BriefDocComment(BriefDocComment), TypeDistance(TypeDistance) {
     assert(CompletionString);
     AssociatedKind = static_cast<unsigned>(Kind);
     IsSystem = 0;

include/swift/Sema/IDETypeChecking.h

@@ -22,6 +22,7 @@
 #include "swift/AST/Identifier.h"
 #include "swift/Basic/SourceLoc.h"
 #include <memory>
+#include <tuple>
 
 namespace swift {
   class AbstractFunctionDecl;
@@ -254,6 +255,34 @@ 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);
+
+  /// Compute the insertion location, indentation string, and component type
+  /// for a Fix-It that adds a new build* function to a function builder.
+  std::tuple<SourceLoc, std::string, Type>
+  determineFunctionBuilderBuildFixItInfo(NominalTypeDecl *builder);
 }
 
 #endif

lib/IDE/CodeCompletion.cpp

@@ -1345,13 +1345,17 @@ CodeCompletionResult *CodeCompletionResultBuilder::takeResult() {
 
   case CodeCompletionResult::ResultKind::Keyword:
     return new (*Sink.Allocator)
-        CodeCompletionResult(KeywordKind, SemanticContext, NumBytesToErase,
-                             CCS, ExpectedTypeRelation);
+        CodeCompletionResult(
+          KeywordKind, SemanticContext, NumBytesToErase,
+          CCS, ExpectedTypeRelation,
+          copyString(*Sink.Allocator, BriefDocComment));
 
   case CodeCompletionResult::ResultKind::BuiltinOperator:
   case CodeCompletionResult::ResultKind::Pattern:
     return new (*Sink.Allocator) CodeCompletionResult(
-        Kind, SemanticContext, NumBytesToErase, CCS, ExpectedTypeRelation);
+        Kind, SemanticContext, NumBytesToErase, CCS, ExpectedTypeRelation,
+        CodeCompletionOperatorKind::None,
+        copyString(*Sink.Allocator, BriefDocComment));
 
   case CodeCompletionResult::ResultKind::Literal:
     assert(LiteralKind.hasValue());
@@ -5136,6 +5140,100 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
     }
   }
 
+  static StringRef getFunctionBuilderDocComment(
+      FunctionBuilderBuildFunction function) {
+    switch (function) {
+    case FunctionBuilderBuildFunction::BuildArray:
+      return "Enables support for..in loops in a function builder by "
+        "combining the results of all iterations into a single result";
+
+    case FunctionBuilderBuildFunction::BuildBlock:
+      return "Required by every function builder to build combined results "
+          "from statement blocks";
+
+    case FunctionBuilderBuildFunction::BuildEitherFirst:
+      return "With buildEither(second:), enables support for 'if-else' and "
+          "'switch' statements by folding conditional results into a single "
+          "result";
+
+    case FunctionBuilderBuildFunction::BuildEitherSecond:
+      return "With buildEither(first:), enables support for 'if-else' and "
+          "'switch' statements by folding conditional results into a single "
+          "result";
+
+    case FunctionBuilderBuildFunction::BuildExpression:
+      return "If declared, provides contextual type information for statement "
+          "expressions to translate them into partial results";
+
+    case FunctionBuilderBuildFunction::BuildFinalResult:
+      return "If declared, this will be called on the partial result from the "
+          "outermost block statement to produce the final returned result";
+
+    case FunctionBuilderBuildFunction::BuildLimitedAvailability:
+      return "If declared, this will be called on the partial result of "
+        "an 'if #available' block to allow the function builder to erase "
+        "type information";
+
+    case FunctionBuilderBuildFunction::BuildOptional:
+      return "Enables support for `if` statements that do not have an `else`";
+    }
+  }
+
+  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.addTextChunk("static ");
+
+      Builder.addTextChunk("func ");
+    }
+
+    std::string declStringWithoutFunc;
+    {
+      llvm::raw_string_ostream out(declStringWithoutFunc);
+      printFunctionBuilderBuildFunction(
+          builder, componentType, function, None, out);
+    }
+    Builder.addTextChunk(declStringWithoutFunc);
+    Builder.addBraceStmtWithCursor();
+    Builder.setBriefDocComment(getFunctionBuilderDocComment(function));
+  }
+
+  /// 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 +5252,10 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
       addDesignatedInitializers(NTD);
       addAssociatedTypes(NTD);
     }
+
+    if (NTD && NTD->getAttrs().hasAttribute<FunctionBuilderAttr>()) {
+      addFunctionBuilderBuildCompletions(NTD);
+    }
   }
 };
 } // end anonymous namespace

lib/IDE/CodeCompletionResultBuilder.h

@@ -92,6 +92,7 @@ class CodeCompletionResultBuilder {
   bool IsNotRecommended = false;
   CodeCompletionResult::NotRecommendedReason NotRecReason =
     CodeCompletionResult::NotRecommendedReason::NoReason;
+  StringRef BriefDocComment = StringRef();
 
   void addChunkWithText(CodeCompletionString::Chunk::ChunkKind Kind,
                         StringRef Text);
@@ -446,6 +447,10 @@ class CodeCompletionResultBuilder {
     addWhitespace(space);
     getLastChunk().setIsAnnotation();
   }
+
+  void setBriefDocComment(StringRef comment) {
+    BriefDocComment = comment;
+  }
 };
 
 } // namespace ide