Repository: apple/swift
File: lib/Sema/TypeCheckAttr.cpp (C++)

Code before:

  auto *returnStmt = ReturnStmt::createImplicit(context, returnedExpr);

  SmallVector<ASTNode, 1> stmts;
  stmts.push_back(returnStmt);
  auto *body = BraceStmt::create(context, SourceLoc(), stmts,
                                SourceLoc(), /*Implicit*/true);

  return std::make_pair(body, /*typechecked=*/false);
}

FuncDecl *
SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
                                        Decl *D) const {
  auto &context = D->getASTContext();

  MainTypeAttr *attr = D->getAttrs().getAttribute<MainTypeAttr>();
  if (attr == nullptr)
    return nullptr;

  auto *extension = dyn_cast<ExtensionDecl>(D);

  IterableDeclContext *iterableDeclContext;
  DeclContext *declContext;
  NominalTypeDecl *nominal;
  SourceRange braces;

  if (extension) {
    nominal = extension->getExtendedNominal();
    iterableDeclContext = extension;
    declContext = extension;
    braces = extension->getBraces();
  } else {
    nominal = dyn_cast<NominalTypeDecl>(D);
    iterableDeclContext = nominal;
    declContext = nominal;
    braces = nominal->getBraces();
  }

  assert(nominal && "Should have already recognized that the MainType decl "
                    "isn't applicable to decls other than NominalTypeDecls");
  assert(iterableDeclContext);
  assert(declContext);

  // The type cannot be generic.
  if (nominal->isGenericContext()) {
    context.Diags.diagnose(attr->getLocation(),
                           diag::attr_generic_ApplicationMain_not_supported, 2);
    attr->setInvalid();
    return nullptr;
  }

Review comment (suggestion):
```suggestion
llvm::SmallString<128> generateMainFunctionText(ASTContext &C, NominalTypeDecl *parentDecl,
```

Code after:
    if (mainFunction->hasThrows()) {
      // $main() async throws { try await main() }
      awaitExpr =
          new (context) TryExpr(callExpr->getLoc(), awaitExpr,
                                context.TheEmptyTupleType, /*implicit*/ true);
    }
    returnedExpr = awaitExpr;
  } else if (mainFunction->hasThrows()) {
    auto *tryExpr = new (context) TryExpr(
        callExpr->getLoc(), callExpr, context.TheEmptyTupleType, /*implicit=*/true);
    returnedExpr = tryExpr;
  } else {
    returnedExpr = callExpr;
  }

  auto *returnStmt = ReturnStmt::createImplicit(context, returnedExpr);

  SmallVector<ASTNode, 1> stmts;
  stmts.push_back(returnStmt);
  auto *body = BraceStmt::create(context, SourceLoc(), stmts,
                                SourceLoc(), /*Implicit*/true);

  return std::make_pair(body, /*typechecked=*/false);
}

static llvm::SmallString<128>
generateMainFunctionText(ASTContext &C, NominalTypeDecl *parentDecl,
                         bool isThrows, bool isAsync) {
  StringRef ExtraIndent;
  StringRef CurrentIndent = Lexer::getIndentationForLine(
      C.SourceMgr, parentDecl->getStartLoc(), &ExtraIndent);
  std::string MethodIndent = (CurrentIndent + ExtraIndent).str();

  llvm::SmallString<128> Text;
  llvm::raw_svector_ostream OS(Text);
  ExtraIndentStreamPrinter Printer(OS, MethodIndent);

  Printer.printNewline();

  Printer << "static func main() ";
  if (isAsync)
    Printer << "async ";
  if (isThrows)
    Printer << "throws ";

  // Print the "{ <#code#> }" placeholder body.
  Printer << "{\n";
  Printer.printIndent();
  Printer << ExtraIndent << getCodePlaceholder();
  Printer.printNewline();
  Printer << "}\n";