Merge pull request swiftlang#32767 from slavapestov/requestify-main-attr

Request-ify synthesis of the main function for the @main attribute

Author: slavapestov

include/swift/AST/ASTTypeIDZone.def

@@ -40,6 +40,7 @@ SWIFT_TYPEID_NAMED(ConstructorDecl *, ConstructorDecl)
 SWIFT_TYPEID_NAMED(CustomAttr *, CustomAttr)
 SWIFT_TYPEID_NAMED(Decl *, Decl)
 SWIFT_TYPEID_NAMED(EnumDecl *, EnumDecl)
+SWIFT_TYPEID_NAMED(FuncDecl *, FuncDecl)
 SWIFT_TYPEID_NAMED(GenericParamList *, GenericParamList)
 SWIFT_TYPEID_NAMED(GenericTypeParamType *, GenericTypeParamType)
 SWIFT_TYPEID_NAMED(InfixOperatorDecl *, InfixOperatorDecl)

include/swift/AST/ASTTypeIDs.h

@@ -30,6 +30,7 @@ class ConstructorDecl;
 class CustomAttr;
 class Decl;
 class EnumDecl;
+class FuncDecl;
 enum class FunctionBuilderBodyPreCheck : uint8_t;
 class GenericParamList;
 class GenericSignature;

include/swift/AST/TypeCheckRequests.h

@@ -2564,6 +2564,23 @@ class CustomAttrTypeRequest
   void cacheResult(Type value) const;
 };
 
+class SynthesizeMainFunctionRequest
+    : public SimpleRequest<SynthesizeMainFunctionRequest,
+                           FuncDecl *(Decl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  FuncDecl *evaluate(Evaluator &evaluator, Decl *) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 // Allow AnyValue to compare two Type values, even though Type doesn't
 // support ==.
 template<>

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -274,3 +274,5 @@ SWIFT_REQUEST(TypeChecker, LookupAllConformancesInContextRequest,
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SimpleDidSetRequest, 
               bool(AccessorDecl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, SynthesizeMainFunctionRequest,
+              FuncDecl *(Decl *), Cached, NoLocationInfo)

lib/Sema/TypeCheckAttr.cpp

@@ -1782,7 +1782,76 @@ void AttributeChecker::visitUIApplicationMainAttr(UIApplicationMainAttr *attr) {
                                 C.getIdentifier("UIApplicationMain"));
 }
 
-void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
+namespace {
+struct MainTypeAttrParams {
+  FuncDecl *mainFunction;
+  MainTypeAttr *attr;
+};
+
+}
+static std::pair<BraceStmt *, bool>
+synthesizeMainBody(AbstractFunctionDecl *fn, void *arg) {
+  ASTContext &context = fn->getASTContext();
+  MainTypeAttrParams *params = (MainTypeAttrParams *) arg;
+
+  FuncDecl *mainFunction = params->mainFunction;
+  auto location = params->attr->getLocation();
+  NominalTypeDecl *nominal = fn->getDeclContext()->getSelfNominalTypeDecl();
+
+  auto *typeExpr = TypeExpr::createImplicit(nominal->getDeclaredType(), context);
+
+  SubstitutionMap substitutionMap;
+  if (auto *environment = mainFunction->getGenericEnvironment()) {
+    substitutionMap = SubstitutionMap::get(
+      environment->getGenericSignature(),
+      [&](SubstitutableType *type) { return nominal->getDeclaredType(); },
+      LookUpConformanceInModule(nominal->getModuleContext()));
+  } else {
+    substitutionMap = SubstitutionMap();
+  }
+
+  auto funcDeclRef = ConcreteDeclRef(mainFunction, substitutionMap);
+
+  auto *memberRefExpr = new (context) MemberRefExpr(
+      typeExpr, SourceLoc(), funcDeclRef, DeclNameLoc(location),
+      /*Implicit*/ true);
+  memberRefExpr->setImplicit(true);
+
+  auto *callExpr = CallExpr::createImplicit(context, memberRefExpr, {}, {});
+  callExpr->setImplicit(true);
+  callExpr->setThrows(mainFunction->hasThrows());
+  callExpr->setType(context.TheEmptyTupleType);
+
+  Expr *returnedExpr;
+
+  if (mainFunction->hasThrows()) {
+    auto *tryExpr = new (context) TryExpr(
+        SourceLoc(), callExpr, context.TheEmptyTupleType, /*implicit=*/true);
+    returnedExpr = tryExpr;
+  } else {
+    returnedExpr = callExpr;
+  }
+
+  auto *returnStmt =
+      new (context) ReturnStmt(SourceLoc(), callExpr, /*Implicit=*/true);
+
+  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;
@@ -1802,25 +1871,19 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
     braces = nominal->getBraces();
   }
 
-  if (!nominal) {
-    assert(false && "Should have already recognized that the MainType decl "
+  assert(nominal && "Should have already recognized that the MainType decl "
                     "isn't applicable to decls other than NominalTypeDecls");
-    return;
-  }
   assert(iterableDeclContext);
   assert(declContext);
 
   // The type cannot be generic.
   if (nominal->isGenericContext()) {
-    diagnose(attr->getLocation(),
-             diag::attr_generic_ApplicationMain_not_supported, 2);
+    context.Diags.diagnose(attr->getLocation(),
+                           diag::attr_generic_ApplicationMain_not_supported, 2);
     attr->setInvalid();
-    return;
+    return nullptr;
   }
 
-  SourceFile *file = cast<SourceFile>(declContext->getModuleScopeContext());
-  assert(file);
-
   // Create a function
   //
   //     func $main() {
@@ -1832,8 +1895,6 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
   // usual type-checking.  The alternative would be to directly call
   // mainType.main() from the entry point, and that would require fully
   // type-checking the call to mainType.main().
-  auto &context = D->getASTContext();
-  auto location = attr->getLocation();
 
   auto resolution = resolveValueMember(
       *declContext, nominal->getInterfaceType(), context.Id_main);
@@ -1861,107 +1922,52 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
     }
 
     if (viableCandidates.size() != 1) {
-      diagnose(attr->getLocation(), diag::attr_MainType_without_main,
-               nominal->getBaseName());
+      context.Diags.diagnose(attr->getLocation(),
+                             diag::attr_MainType_without_main,
+                             nominal->getBaseName());
       attr->setInvalid();
-      return;
+      return nullptr;
     }
     mainFunction = viableCandidates[0];
   }
 
-  bool mainFunctionThrows = mainFunction->hasThrows();
-
-  auto voidToVoidFunctionType =
-      FunctionType::get({}, context.TheEmptyTupleType,
-                        FunctionType::ExtInfo().withThrows(mainFunctionThrows));
-  auto nominalToVoidToVoidFunctionType = FunctionType::get({AnyFunctionType::Param(nominal->getInterfaceType())}, voidToVoidFunctionType);
   auto *func = FuncDecl::create(
       context, /*StaticLoc*/ SourceLoc(), StaticSpellingKind::KeywordStatic,
       /*FuncLoc*/ SourceLoc(),
       DeclName(context, DeclBaseName(context.Id_MainEntryPoint),
                ParameterList::createEmpty(context)),
-      /*NameLoc*/ SourceLoc(), /*Throws=*/mainFunctionThrows,
+      /*NameLoc*/ SourceLoc(), /*Throws=*/mainFunction->hasThrows(),
       /*ThrowsLoc=*/SourceLoc(),
       /*GenericParams=*/nullptr, ParameterList::createEmpty(context),
       /*FnRetType=*/TypeLoc::withoutLoc(TupleType::getEmpty(context)),
       declContext);
   func->setImplicit(true);
   func->setSynthesized(true);
 
-  auto *typeExpr = TypeExpr::createImplicit(nominal->getDeclaredType(), context);
-
-  SubstitutionMap substitutionMap;
-  if (auto *environment = mainFunction->getGenericEnvironment()) {
-    substitutionMap = SubstitutionMap::get(
-      environment->getGenericSignature(),
-      [&](SubstitutableType *type) { return nominal->getDeclaredType(); },
-      LookUpConformanceInModule(nominal->getModuleContext()));
-  } else {
-    substitutionMap = SubstitutionMap();
-  }
-
-  auto funcDeclRef = ConcreteDeclRef(mainFunction, substitutionMap);
-
-  auto *memberRefExpr = new (context) MemberRefExpr(
-      typeExpr, SourceLoc(), funcDeclRef, DeclNameLoc(location),
-      /*Implicit*/ true);
-  memberRefExpr->setImplicit(true);
-
-  auto *callExpr = CallExpr::createImplicit(context, memberRefExpr, {}, {});
-  callExpr->setImplicit(true);
-  callExpr->setThrows(mainFunctionThrows);
-  callExpr->setType(context.TheEmptyTupleType);
-
-  Expr *returnedExpr;
+  auto *params = context.Allocate<MainTypeAttrParams>();
+  params->mainFunction = mainFunction;
+  params->attr = attr;
+  func->setBodySynthesizer(synthesizeMainBody, params);
 
-  if (mainFunctionThrows) {
-    auto *tryExpr = new (context) TryExpr(
-        SourceLoc(), callExpr, context.TheEmptyTupleType, /*implicit=*/true);
-    returnedExpr = tryExpr;
-  } else {
-    returnedExpr = callExpr;
-  }
+  iterableDeclContext->addMember(func);
 
-  auto *returnStmt =
-      new (context) ReturnStmt(SourceLoc(), callExpr, /*Implicit=*/true);
+  return func;
+}
 
-  SmallVector<ASTNode, 1> stmts;
-  stmts.push_back(returnStmt);
-  auto *body = BraceStmt::create(context, SourceLoc(), stmts,
-                                SourceLoc(), /*Implicit*/true);
-  func->setBodyParsed(body);
-  func->setInterfaceType(nominalToVoidToVoidFunctionType);
+void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
+  auto &context = D->getASTContext();
 
-  iterableDeclContext->addMember(func);
+  SourceFile *file = D->getDeclContext()->getParentSourceFile();
+  assert(file);
 
-  // This function must be type-checked. Why?  Consider the following scenario:
-  //
-  //     protocol AlmostMainable {}
-  //     protocol ReallyMainable {}
-  //     extension AlmostMainable where Self : ReallyMainable {
-  //         static func main() {}
-  //     }
-  //     @main struct Main : AlmostMainable {}
-  //
-  // Note in particular that Main does not conform to ReallyMainable.
-  //
-  // In this case, resolveValueMember will find the function main in the 
-  // extension, and so, since there is one candidate, the function $main will
-  // accordingly be formed as usual:
-  //
-  //     func $main() {
-  //         return Main.main()
-  //     }
-  //
-  // Of course, this function's body does not type-check.
-  file->DelayedFunctions.push_back(func);
+  auto *func = evaluateOrDefault(context.evaluator,
+                                 SynthesizeMainFunctionRequest{D},
+                                 nullptr);
 
   // Register the func as the main decl in the module. If there are multiples
   // they will be diagnosed.
-  if (file->registerMainDecl(func, attr->getLocation())) {
+  if (file->registerMainDecl(func, attr->getLocation()))
     attr->setInvalid();
-    return;
-  }
 }
 
 /// Determine whether the given context is an extension to an Objective-C class

lib/Sema/TypeCheckDecl.cpp

@@ -2523,6 +2523,12 @@ EmittedMembersRequest::evaluate(Evaluator &evaluator,
   forceConformance(Context.getProtocol(KnownProtocolKind::Hashable));
   forceConformance(Context.getProtocol(KnownProtocolKind::Differentiable));
 
+  // If the class has a @main attribute, we need to force synthesis of the
+  // $main function.
+  (void) evaluateOrDefault(Context.evaluator,
+                           SynthesizeMainFunctionRequest{CD},
+                           nullptr);
+
   for (auto *member : CD->getMembers()) {
     if (auto *var = dyn_cast<VarDecl>(member)) {
       // The projected storage wrapper ($foo) might have dynamically-dispatched

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -1797,13 +1797,13 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     checkGenericParams(ED);
 
     // Check for circular inheritance of the raw type.
-    (void)ED->hasCircularRawValue();
+    (void) ED->hasCircularRawValue();
+
+    TypeChecker::checkDeclAttributes(ED);
 
     for (Decl *member : ED->getMembers())
       visit(member);
 
-    TypeChecker::checkDeclAttributes(ED);
-
     checkInheritanceClause(ED);
 
     checkAccessControl(ED);
@@ -1845,13 +1845,13 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
 
     installCodingKeysIfNecessary(SD);
 
+    TypeChecker::checkDeclAttributes(SD);
+
     for (Decl *Member : SD->getMembers())
       visit(Member);
 
     TypeChecker::checkPatternBindingCaptures(SD);
 
-    TypeChecker::checkDeclAttributes(SD);
-
     checkInheritanceClause(SD);
 
     checkAccessControl(SD);
@@ -1974,6 +1974,8 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     // Force creation of an implicit destructor, if any.
     (void) CD->getDestructor();
 
+    TypeChecker::checkDeclAttributes(CD);
+
     for (Decl *Member : CD->getEmittedMembers())
       visit(Member);
 
@@ -2084,8 +2086,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
       }
     }
 
-    TypeChecker::checkDeclAttributes(CD);
-
     checkInheritanceClause(CD);
 
     checkAccessControl(CD);
@@ -2105,12 +2105,12 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     // Check for circular inheritance within the protocol.
     (void)PD->hasCircularInheritedProtocols();
 
+    TypeChecker::checkDeclAttributes(PD);
+
     // Check the members.
     for (auto Member : PD->getMembers())
       visit(Member);
 
-    TypeChecker::checkDeclAttributes(PD);
-
     checkAccessControl(PD);
 
     checkInheritanceClause(PD);
@@ -2428,14 +2428,15 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
 
     checkGenericParams(ED);
 
+    TypeChecker::checkDeclAttributes(ED);
+
     for (Decl *Member : ED->getMembers())
       visit(Member);
 
     TypeChecker::checkPatternBindingCaptures(ED);
 
     TypeChecker::checkConformancesInContext(ED);
 
-    TypeChecker::checkDeclAttributes(ED);
     checkAccessControl(ED);
 
     checkExplicitAvailability(ED);