[Function builders] Add stubs for missing buildBlock and buildLimitedAvailability

The lack of these functions in a function builder is diagnosed in
different places in the type checker. Extend these diagnostics to
provide stubs in the function builder.

Thanks to Suyash for the suggestion!

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -5321,6 +5321,9 @@ NOTE(function_builder_missing_build_either, none,
 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/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;
@@ -277,6 +278,11 @@ namespace swift {
       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/Sema/BuilderTransform.cpp

@@ -1247,6 +1247,32 @@ class BuilderClosureRewriter
           ctx.Diags.diagnose(
               loc, diag::function_builder_missing_limited_availability,
               builderTransform.builderType);
+
+          // Add a note to the function builder with a stub for
+          // buildLimitedAvailability().
+          if (auto builder = builderTransform.builderType->getAnyNominal()) {
+            SourceLoc buildInsertionLoc;
+            std::string stubIndent;
+            Type componentType;
+            std::tie(buildInsertionLoc, stubIndent, componentType) =
+                determineFunctionBuilderBuildFixItInfo(builder);
+            if (buildInsertionLoc.isValid()) {
+              std::string fixItString;
+              {
+                llvm::raw_string_ostream out(fixItString);
+                printFunctionBuilderBuildFunction(
+                    builder, componentType,
+                    FunctionBuilderBuildFunction::BuildLimitedAvailability,
+                    stubIndent, out);
+
+                builder->diagnose(
+                    diag::function_builder_missing_build_limited_availability,
+                    builderTransform.builderType)
+                  .fixItInsert(buildInsertionLoc, fixItString);
+              }
+            }
+          }
+
           return true;
         }
 
@@ -1917,6 +1943,28 @@ Type swift::inferFunctionBuilderComponentType(NominalTypeDecl *builder) {
   return componentType;
 }
 
+std::tuple<SourceLoc, std::string, Type>
+swift::determineFunctionBuilderBuildFixItInfo(NominalTypeDecl *builder) {
+  SourceLoc buildInsertionLoc = builder->getBraces().Start;
+  std::string stubIndent;
+  Type componentType;
+
+  if (buildInsertionLoc.isInvalid())
+    return std::make_tuple(buildInsertionLoc, stubIndent, componentType);
+
+  ASTContext &ctx = builder->getASTContext();
+  buildInsertionLoc = Lexer::getLocForEndOfToken(
+      ctx.SourceMgr, buildInsertionLoc);
+
+  StringRef extraIndent;
+  StringRef currentIndent = Lexer::getIndentationForLine(
+      ctx.SourceMgr, buildInsertionLoc, &extraIndent);
+  stubIndent = (currentIndent + extraIndent).str();
+
+  componentType = inferFunctionBuilderComponentType(builder);
+  return std::make_tuple(buildInsertionLoc, stubIndent, componentType);
+}
+
 void swift::printFunctionBuilderBuildFunction(
       NominalTypeDecl *builder, Type componentType,
       FunctionBuilderBuildFunction function,

lib/Sema/CSDiagnostics.cpp

@@ -5511,21 +5511,11 @@ void SkipUnhandledConstructInFunctionBuilderFailure::diagnosePrimary(
 
     // Emit custom notes to help the user introduce the appropriate 'build'
     // functions.
-    SourceLoc buildInsertionLoc = builder->getBraces().Start;
+    SourceLoc buildInsertionLoc;
     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);
-    }
+    std::tie(buildInsertionLoc, stubIndent, componentType) =
+        determineFunctionBuilderBuildFixItInfo(builder);
 
     if (buildInsertionLoc.isInvalid()) {
       // Do nothing.

lib/Sema/TypeCheckAttr.cpp

@@ -3024,7 +3024,29 @@ void AttributeChecker::visitFunctionBuilderAttr(FunctionBuilderAttr *attr) {
       /*argLabels=*/{}, &potentialMatches);
 
   if (!supportsBuildBlock) {
-    diagnose(nominal->getLoc(), diag::function_builder_static_buildblock);
+    {
+      auto diag = diagnose(
+          nominal->getLoc(), diag::function_builder_static_buildblock);
+
+      // If there were no close matches, propose adding a stub.
+      SourceLoc buildInsertionLoc;
+      std::string stubIndent;
+      Type componentType;
+      std::tie(buildInsertionLoc, stubIndent, componentType) =
+          determineFunctionBuilderBuildFixItInfo(nominal);
+      if (buildInsertionLoc.isValid() && potentialMatches.empty()) {
+        std::string fixItString;
+        {
+          llvm::raw_string_ostream out(fixItString);
+          printFunctionBuilderBuildFunction(
+              nominal, componentType,
+              FunctionBuilderBuildFunction::BuildBlock,
+              stubIndent, out);
+        }
+
+        diag.fixItInsert(buildInsertionLoc, fixItString);
+      }
+    }
 
     // For any close matches, attempt to explain to the user why they aren't
     // valid.

test/Constraints/function_builder_availability.swift

@@ -8,7 +8,7 @@ enum Either<T,U> {
 }
 
 @_functionBuilder
-struct TupleBuilder {
+struct TupleBuilder { // expected-note{{add 'buildLimitedAvailability(_:)' to the function builder 'TupleBuilder' to erase type information for less-available types}}{{22-22=\n    static func buildLimitedAvailability(_ component: <#Component#>) -> <#Component#> {\n      <#code#>\n    \}}}
   static func buildBlock<T1>(_ t1: T1) -> (T1) {
     return (t1)
   }

test/decl/var/function_builders.swift

@@ -7,7 +7,7 @@ var globalBuilder: Int
 func globalBuilderFunction() -> Int { return 0 }
 
 @_functionBuilder
-struct Maker {} // expected-error {{function builder must provide at least one static 'buildBlock' method}}
+struct Maker {} // expected-error {{function builder must provide at least one static 'buildBlock' method}}{{15-15=\n    static func buildBlock(_ components: <#Component#>...) -> <#Component#> {\n      <#code#>\n    \}}}
 
 @_functionBuilder
 class Inventor {} // expected-error {{function builder must provide at least one static 'buildBlock' method}}