Sema: Lazily synthesize body of main function

Author: slavapestov

lib/Sema/TypeCheckAttr.cpp

@@ -1782,6 +1782,67 @@ void AttributeChecker::visitUIApplicationMainAttr(UIApplicationMainAttr *attr) {
                                 C.getIdentifier("UIApplicationMain"));
 }
 
+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);
+}
+
 void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
   auto *extension = dyn_cast<ExtensionDecl>(D);
 
@@ -1802,11 +1863,8 @@ 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);
 
@@ -1833,7 +1891,6 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
   // 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);
@@ -1869,63 +1926,23 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
     mainFunction = viableCandidates[0];
   }
 
-  bool mainFunctionThrows = mainFunction->hasThrows();
-
   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;
-
-  if (mainFunctionThrows) {
-    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);
-  func->setBodyParsed(body);
+  auto *params = context.Allocate<MainTypeAttrParams>();
+  params->mainFunction = mainFunction;
+  params->attr = attr;
+  func->setBodySynthesizer(synthesizeMainBody, params);
 
   iterableDeclContext->addMember(func);