@main: Allowed attribute on extensions.

Previously the @main attribute could only be applied to
NominalTypeDecls.  Here, that limitation is lifted so that the attribute
can be applied to ExtensionDecls.

Author: nate-chandler

include/swift/AST/Attr.def

@@ -148,7 +148,7 @@ SIMPLE_DECL_ATTR(dynamicCallable, DynamicCallable,
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   6)
 SIMPLE_DECL_ATTR(main, MainType,
-  OnClass | OnStruct | OnEnum | 
+  OnClass | OnStruct | OnEnum | OnExtension |
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   7)
 SIMPLE_DECL_ATTR(_exported, Exported,

lib/SILGen/SILGenFunction.cpp

@@ -732,7 +732,15 @@ void SILGenFunction::emitArtificialTopLevel(Decl *mainDecl) {
     SILFunction *mainFunction =
         SGM.getFunction(mainFunctionDeclRef, NotForDefinition);
 
-    NominalTypeDecl *mainType = cast<NominalTypeDecl>(mainFunc->getDeclContext());
+    ExtensionDecl *mainExtension =
+        dyn_cast<ExtensionDecl>(mainFunc->getDeclContext());
+
+    NominalTypeDecl *mainType;
+    if (mainExtension) {
+      mainType = mainExtension->getExtendedNominal();
+    } else {
+      mainType = cast<NominalTypeDecl>(mainFunc->getDeclContext());
+    }
     auto metatype = B.createMetatype(mainType, getLoweredType(mainType->getInterfaceType()));
 
     auto mainFunctionRef = B.createFunctionRef(mainFunc, mainFunction);

lib/Sema/TypeCheckAttr.cpp

@@ -1700,13 +1700,32 @@ void AttributeChecker::visitUIApplicationMainAttr(UIApplicationMainAttr *attr) {
 }
 
 void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
-  auto *nominal = dyn_cast<NominalTypeDecl>(D);
+  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();
+  }
 
   if (!nominal) {
     assert(false && "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()) {
@@ -1716,7 +1735,8 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
     return;
   }
 
-  auto *file = cast<SourceFile>(nominal->getModuleScopeContext());
+  SourceFile *file = cast<SourceFile>(declContext->getModuleScopeContext());
+  assert(file);
 
   // Create a function
   //
@@ -1729,11 +1749,11 @@ 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 = nominal->getASTContext();
+  auto &context = D->getASTContext();
   auto location = attr->getLocation();
 
-  auto resolution = resolveValueMember(*nominal, nominal->getInterfaceType(),
-                                       context.Id_main);
+  auto resolution = resolveValueMember(
+      *declContext, nominal->getInterfaceType(), context.Id_main);
 
   FuncDecl *mainFunction = nullptr;
 
@@ -1769,14 +1789,14 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
   auto voidToVoidFunctionType = FunctionType::get({}, context.TheEmptyTupleType);
   auto nominalToVoidToVoidFunctionType = FunctionType::get({AnyFunctionType::Param(nominal->getInterfaceType())}, voidToVoidFunctionType);
   auto *func = FuncDecl::create(
-      context, /*StaticLoc*/ nominal->getBraces().End, StaticSpellingKind::KeywordStatic,
+      context, /*StaticLoc*/ braces.End, StaticSpellingKind::KeywordStatic,
       /*FuncLoc*/ location,
       DeclName(context, DeclBaseName(context.Id_MainEntryPoint),
                ParameterList::createEmpty(context)),
-      /*NameLoc*/ nominal->getBraces().End, /*Throws=*/false, /*ThrowsLoc=*/SourceLoc(),
+      /*NameLoc*/ braces.End, /*Throws=*/false, /*ThrowsLoc=*/SourceLoc(),
       /*GenericParams=*/nullptr, ParameterList::createEmpty(context),
       /*FnRetType=*/TypeLoc::withoutLoc(TupleType::getEmpty(context)),
-      nominal);
+      declContext);
   func->setImplicit(true);
   func->setSynthesized(true);
 
@@ -1817,7 +1837,8 @@ void AttributeChecker::visitMainTypeAttr(MainTypeAttr *attr) {
   func->setBodyParsed(body);
   func->setInterfaceType(nominalToVoidToVoidFunctionType);
 
-  nominal->addMember(func);
+  iterableDeclContext->addMember(func);
+
   // This function must be type-checked. Why?  Consider the following scenario:
   //
   //     protocol AlmostMainable {}

test/attr/ApplicationMain/attr_main_class_extension.swift

@@ -0,0 +1,12 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+class EntryPoint {
+}
+
+@main
+extension EntryPoint {
+  static func main() {
+  }
+}
+
+

test/attr/ApplicationMain/attr_main_class_extension_external_main.swift

@@ -0,0 +1,11 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+class EntryPoint {
+}
+
+@main
+extension EntryPoint {
+  static func main() {
+  }
+}
+

test/attr/ApplicationMain/attr_main_class_extension_multi_module/A.swift

@@ -0,0 +1,6 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+public class Main {
+}

test/attr/ApplicationMain/attr_main_class_extension_multi_module/main.swift

@@ -0,0 +1,12 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module -emit-library -module-name ModuleA -module-link-name ModuleA %S/A.swift -o %t/%target-library-name(ModuleA)
+// RUN: %target-swift-frontend -c -I %t -L %t -lModuleA -parse-as-library %s 
+
+import ModuleA
+
+@main
+extension Main {
+  static func main() {
+    print("ok")
+  }
+}

test/attr/ApplicationMain/attr_main_class_extension_multi_module_external_main/A.swift

@@ -0,0 +1,9 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+public class Main {
+  public class func main() {
+    print("ok")
+  }
+}

test/attr/ApplicationMain/attr_main_class_extension_multi_module_external_main/main.swift

@@ -0,0 +1,9 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module -emit-library -module-name ModuleA -module-link-name ModuleA %S/A.swift -o %t/%target-library-name(ModuleA)
+// RUN: %target-swift-frontend -c -I %t -L %t -lModuleA -parse-as-library %s 
+
+import ModuleA
+
+@main
+extension Main {
+}

test/attr/ApplicationMain/attr_main_enum_extension.swift

@@ -0,0 +1,13 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+enum EntryPoint {
+}
+
+@main
+extension EntryPoint {
+  static func main() {
+  }
+}
+
+
+

test/attr/ApplicationMain/attr_main_enum_extension_external_main.swift

@@ -0,0 +1,13 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+enum EntryPoint {
+  static func main() {
+  }
+}
+
+@main
+extension EntryPoint {
+}
+
+
+

test/attr/ApplicationMain/attr_main_enum_extension_multi_module/A.swift

@@ -0,0 +1,6 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+public enum Main {
+}

test/attr/ApplicationMain/attr_main_enum_extension_multi_module/main.swift

@@ -0,0 +1,12 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module -emit-library -module-name ModuleA -module-link-name ModuleA %S/A.swift -o %t/%target-library-name(ModuleA)
+// RUN: %target-swift-frontend -c -I %t -L %t -lModuleA -parse-as-library %s 
+
+import ModuleA
+
+@main
+extension Main {
+  static func main() {
+    print("ok")
+  }
+}

test/attr/ApplicationMain/attr_main_enum_extension_multi_module_external_main/A.swift

@@ -0,0 +1,9 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+public enum Main {
+  public static func main() {
+    print("ok")
+  }
+}

test/attr/ApplicationMain/attr_main_enum_extension_multi_module_external_main/main.swift

@@ -0,0 +1,9 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module -emit-library -module-name ModuleA -module-link-name ModuleA %S/A.swift -o %t/%target-library-name(ModuleA)
+// RUN: %target-swift-frontend -c -I %t -L %t -lModuleA -parse-as-library %s 
+
+import ModuleA
+
+@main
+extension Main {
+}

test/attr/ApplicationMain/attr_main_extension_multi_file/main1.swift

@@ -0,0 +1,10 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+@main
+extension Main {
+  static func main() {
+    print("hello world")
+  }
+}

test/attr/ApplicationMain/attr_main_extension_multi_file/main2.swift

@@ -0,0 +1,11 @@
+// RUN: %target-swift-frontend -typecheck -verify %s %S/main1.swift
+
+// Serialized partial AST support:
+// RUN: %target-swift-frontend -module-name main -emit-module-path %t.swiftmodule -primary-file %s %S/main1.swift
+// RUN: %target-swift-frontend -module-name main -parse-as-library -typecheck %t.swiftmodule -primary-file %S/main1.swift -verify -verify-ignore-unknown
+
+class Main {
+}
+
+// FIXME: Remove -verify-ignore-unknown.
+// <unknown>:0: error: unexpected error produced: 'main' attribute can only apply to one class in a module

test/attr/ApplicationMain/attr_main_extension_nofunc.swift

@@ -0,0 +1,10 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+class EntryPoint {
+}
+
+@main // expected-error{{'EntryPoint' is annotated with @main and must provide a main static function of type () -> ()}}
+extension EntryPoint {
+}
+
+

test/attr/ApplicationMain/attr_main_struct_extension.swift

@@ -0,0 +1,11 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+struct EntryPoint {
+}
+
+@main
+extension EntryPoint {
+  static func main() {
+  }
+}
+

test/attr/ApplicationMain/attr_main_struct_extension_external_main.swift

@@ -0,0 +1,11 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+struct EntryPoint {
+  static func main() {
+  }
+}
+
+@main
+extension EntryPoint {
+}
+

test/attr/ApplicationMain/attr_main_struct_extension_multi_module/A.swift

@@ -0,0 +1,6 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+public struct Main {
+}

test/attr/ApplicationMain/attr_main_struct_extension_multi_module/main.swift

@@ -0,0 +1,12 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module -emit-library -module-name ModuleA -module-link-name ModuleA %S/A.swift -o %t/%target-library-name(ModuleA)
+// RUN: %target-swift-frontend -c -I %t -L %t -lModuleA -parse-as-library %s 
+
+import ModuleA
+
+@main
+extension Main {
+  static func main() {
+    print("ok")
+  }
+}

test/attr/ApplicationMain/attr_main_struct_extension_multi_module_external_main/A.swift

@@ -0,0 +1,9 @@
+// This file is a part of the multi-file test driven by 'main2.swift'.
+
+// RUN: %target-swift-frontend -parse %s
+
+public struct Main {
+  public static func main() {
+    print("ok")
+  }
+}

test/attr/ApplicationMain/attr_main_struct_extension_multi_module_external_main/main.swift

@@ -0,0 +1,9 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module -emit-library -module-name ModuleA -module-link-name ModuleA %S/A.swift -o %t/%target-library-name(ModuleA)
+// RUN: %target-swift-frontend -c -I %t -L %t -lModuleA -parse-as-library %s 
+
+import ModuleA
+
+@main
+extension Main {
+}

test/attr/ApplicationMain/attr_main_tuple_extension.swift

@@ -0,0 +1,11 @@
+// RUN: %target-swift-frontend -typecheck -parse-as-library -verify %s
+
+@main
+extension (Int, String) { // expected-error {{non-nominal type '(Int, String)' cannot be extended}}
+  static func main() {
+  }
+}
+
+
+
+