[Function builders] Fix-Its and code completion improvements

include/swift/AST/DiagnosticsSema.def

@@ -5312,6 +5312,16 @@ 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))
+
 //------------------------------------------------------------------------------
 // MARK: Tuple Shuffle Diagnostics
 //------------------------------------------------------------------------------

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"
@@ -5490,12 +5491,93 @@ SourceLoc SkipUnhandledConstructInFunctionBuilderFailure::getLoc() const {
   return unhandled.get<Decl *>()->getLoc();
 }
 
+/// Determine whether the given "if" chain has a missing "else".
+static bool hasMissingElseInChain(IfStmt *ifStmt) {
+  if (!ifStmt->getElseStmt())
+    return true;
+
+  if (auto ifElse = dyn_cast<IfStmt>(ifStmt->getElseStmt()))
+    return hasMissingElseInChain(ifElse);
+
+  return false;
+}
+
 void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
     bool asNote) {
-  if (unhandled.is<Stmt *>()) {
+  if (auto stmt = unhandled.dyn_cast<Stmt *>()) {
     emitDiagnostic(asNote ? diag::note_function_builder_control_flow
                           : diag::function_builder_control_flow,
                    builder->getName());
+
+    // 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()) {
+      // Do nothing.
+    } else if (isa<IfStmt>(stmt) && hasMissingElseInChain(cast<IfStmt>(stmt))) {
+      auto diag = emitDiagnosticAt(
+          builder->getLoc(), diag::function_builder_missing_build_optional,
+          builder->getDeclaredInterfaceType());
+
+      std::string fixItString;
+      {
+        llvm::raw_string_ostream out(fixItString);
+        printFunctionBuilderBuildFunction(
+            builder, componentType, FunctionBuilderBuildFunction::BuildOptional,
+            stubIndent, out);
+      }
+
+      diag.fixItInsert(buildInsertionLoc, fixItString);
+    } else if (isa<SwitchStmt>(stmt) || isa<IfStmt>(stmt)) {
+      auto diag = emitDiagnosticAt(
+          builder->getLoc(), diag::function_builder_missing_build_either,
+          builder->getDeclaredInterfaceType());
+
+      std::string fixItString;
+      {
+        llvm::raw_string_ostream out(fixItString);
+        printFunctionBuilderBuildFunction(
+            builder, componentType,
+            FunctionBuilderBuildFunction::BuildEitherFirst,
+            stubIndent, out);
+        out << '\n';
+        printFunctionBuilderBuildFunction(
+            builder, componentType,
+            FunctionBuilderBuildFunction::BuildEitherSecond,
+            stubIndent, out);
+      }
+
+      diag.fixItInsert(buildInsertionLoc, fixItString);
+    } else if (isa<ForEachStmt>(stmt)) {
+      auto diag = emitDiagnosticAt(
+          builder->getLoc(), diag::function_builder_missing_build_array,
+          builder->getDeclaredInterfaceType());
+
+      std::string fixItString;
+      {
+        llvm::raw_string_ostream out(fixItString);
+        printFunctionBuilderBuildFunction(
+            builder, componentType, FunctionBuilderBuildFunction::BuildArray,
+            stubIndent, out);
+      }
+
+      diag.fixItInsert(buildInsertionLoc, fixItString);
+    }
   } else {
     emitDiagnostic(asNote ? diag::note_function_builder_decl
                           : diag::function_builder_decl,

test/Constraints/function_builder_diags.swift

@@ -45,7 +45,10 @@ struct TupleBuilder { // expected-note 2 {{struct 'TupleBuilder' declared here}}
 }
 
 @_functionBuilder
-struct TupleBuilderWithoutIf { // expected-note {{struct 'TupleBuilderWithoutIf' declared here}}
+struct TupleBuilderWithoutIf { // expected-note 3{{struct 'TupleBuilderWithoutIf' declared here}}
+  // expected-note@-1{{add 'buildOptional(_:)' to the function builder 'TupleBuilderWithoutIf' to add support for 'if' statements without an 'else'}}{{31-31=\n    static func buildOptional(_ component: <#Component#>?) -> <#Component#> {\n      <#code#>\n    \}}}
+  // expected-note@-2{{add 'buildEither(first:)' and `buildEither(second:)' to the function builder 'TupleBuilderWithoutIf' to add support for 'if'-'else' and 'switch'}}{{31-31=\n    static func buildEither(first component: <#Component#>) -> <#Component#> {\n      <#code#>\n    \}\n\n    static func buildEither(second component: <#Component#>) -> <#Component#> {\n      <#code#>\n    \}}}
+  // expected-note@-3{{add 'buildArray(_:)' to the function builder 'TupleBuilderWithoutIf' to add support for 'for'..'in' loops}}{{31-31=\n    static func buildArray(_ components: [<#Component#>]) -> <#Component#> {\n      <#code#>\n    \}}}
   static func buildBlock() -> () { }
 
   static func buildBlock<T1>(_ t1: T1) -> T1 {
@@ -106,6 +109,19 @@ func testDiags() {
       "hello"
     }
   }
+
+  tuplifyWithoutIf(true) {
+    if $0 {    // expected-error{{closure containing control flow statement cannot be used with function builder 'TupleBuilderWithoutIf'}}
+      "hello"
+    } else {
+    }
+  }
+
+  tuplifyWithoutIf(true) { a in
+    for x in 0..<100 {    // expected-error{{closure containing control flow statement cannot be used with function builder 'TupleBuilderWithoutIf'}}
+      x
+    }
+  }
 }
 
 struct A { }