[Diagnostics] Add fix-it to @main struct without main static function.

Change-Id: Id23067a3b71eab0927adb18b2c84b575e1dbfa45

Author: CrazyFanFan

lib/Sema/TypeCheckAttr.cpp

@@ -22,6 +22,7 @@
 #include "TypeCheckObjC.h"
 #include "TypeCheckType.h"
 #include "TypeChecker.h"
+#include "swift/AST/ASTPrinter.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/AvailabilityInference.h"
 #include "swift/AST/ClangModuleLoader.h"
@@ -3064,6 +3065,37 @@ synthesizeMainBody(AbstractFunctionDecl *fn, void *arg) {
   return std::make_pair(body, /*typechecked=*/false);
 }
 
+llvm::SmallString<128> generateMainFunctionText(ASTContext &C, Decl *decl,
+                                                bool isThrows, bool isAsync) {
+  // Prepare the indent (same as `printRequirementStub`)
+  StringRef ExtraIndent;
+  StringRef CurrentIndent = Lexer::getIndentationForLine(
+      C.SourceMgr, decl->getStartLoc(), &ExtraIndent);
+
+  llvm::SmallString<128> Text;
+  llvm::raw_svector_ostream OS(Text);
+  ExtraIndentStreamPrinter Printer(OS, CurrentIndent);
+
+  Printer.printNewline();
+  Printer.printIndent();
+
+  Printer << ExtraIndent << "static func main() ";
+  if (isAsync)
+    Printer << "async ";
+  if (isThrows)
+    Printer << "throws ";
+
+  /// Print the "{ <#code#> }" placeholder body
+  Printer << "{\n";
+  Printer.printIndent();
+  Printer << ExtraIndent << ExtraIndent << getCodePlaceholder();
+  Printer.printNewline();
+  Printer.printIndent();
+  Printer << ExtraIndent << "}\n";
+
+  return Text;
+}
+
 FuncDecl *
 SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
                                         Decl *D) const {
@@ -3189,13 +3221,30 @@ SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
         candidates[0].overloadChoices[locator].choice.getDecl());
   }
 
+  auto location = braces.Start;
+
   if (!mainFunction) {
     const bool hasAsyncSupport =
         AvailabilityRange::forDeploymentTarget(context).isContainedIn(
             context.getBackDeployedConcurrencyAvailability());
-    context.Diags.diagnose(attr->getLocation(),
-                           diag::attr_MainType_without_main,
-                           nominal, hasAsyncSupport);
+    auto params =
+        hasAsyncSupport
+            ? std::vector<std::pair<bool, bool>>{{false, false},
+                                                 {true, false},
+                                                 {false, true},
+                                                 {true, true}}
+            : std::vector<std::pair<bool, bool>>{{false, false}, {true, false}};
+
+    auto diag = context.Diags.diagnose(attr->getLocation(),
+                                       diag::attr_MainType_without_main,
+                                       nominal, hasAsyncSupport);
+
+    for (auto [isThrowing, isAsync] : params) {
+      diag.fixItInsertAfter(
+          location,
+          generateMainFunctionText(context, D, isThrowing, isAsync).str());
+    }
+
     attr->setInvalid();
     return nullptr;
   }

test/attr/ApplicationMain/attr_main_instance.swift

@@ -1,6 +1,6 @@
 // RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
 
-@main // expected-error{{'MyBase' is annotated with '@main' and must provide a main static function}}
+@main // expected-error{{'MyBase' is annotated with '@main' and must provide a main static function}} {{4:15-15=\n    static func main() {\n        <#code#>\n    }\n}} {{4:15-15=\n    static func main() throws {\n        <#code#>\n    }\n}} {{4:15-15=\n    static func main() async {\n        <#code#>\n    }\n}} {{4:15-15=\n    static func main() async throws {\n        <#code#>\n    }\n}}
 class MyBase {
   func main() {
   }