[AST] Retrieve @main type from module.

Added a convenience function to ModuleDecl to look up the entry point
@main type if there is one.

Author: nate-chandler

include/swift/AST/Module.h

@@ -841,6 +841,8 @@ class ModuleDecl
     return EntryPointInfo.hasEntryPoint();
   }
 
+  NominalTypeDecl *getMainTypeDecl() const;
+
   /// Returns the associated clang module if one exists.
   const clang::Module *findUnderlyingClangModule() const;
 

lib/AST/Module.cpp

@@ -1876,6 +1876,22 @@ bool SourceFile::registerMainDecl(ValueDecl *mainDecl, SourceLoc diagLoc) {
   return false;
 }
 
+NominalTypeDecl *ModuleDecl::getMainTypeDecl() const {
+  if (!EntryPointInfo.hasEntryPoint())
+    return nullptr;
+  auto *file = EntryPointInfo.getEntryPointFile();
+  if (!file)
+    return nullptr;
+  auto *mainDecl = file->getMainDecl();
+  if (!mainDecl)
+    return nullptr;
+  auto *func = dyn_cast<FuncDecl>(file->getMainDecl());
+  if (!func)
+    return nullptr;
+  auto *nominalType = dyn_cast<NominalTypeDecl>(func->getDeclContext());
+  return nominalType;
+}
+
 bool ModuleDecl::registerEntryPointFile(FileUnit *file, SourceLoc diagLoc,
                                         Optional<ArtificialMainKind> kind) {
   if (!EntryPointInfo.hasEntryPoint()) {