[SILGen] Add SILDeclRef for main entry-point

include/swift/AST/FileUnit.h

@@ -272,7 +272,7 @@ class FileUnit : public DeclContext {
     return dyn_cast_or_null<ClassDecl>(getMainDecl());
   }
   bool hasMainDecl() const { return getMainDecl(); }
-  virtual Decl *getMainDecl() const { return nullptr; }
+  virtual ValueDecl *getMainDecl() const { return nullptr; }
   FuncDecl *getMainFunc() const {
     return dyn_cast_or_null<FuncDecl>(getMainDecl());
   }

include/swift/AST/SourceFile.h

@@ -99,7 +99,7 @@ class SourceFile final : public FileUnit {
 
   /// Either the class marked \@NS/UIApplicationMain or the synthesized FuncDecl
   /// that calls main on the type marked @main.
-  Decl *MainDecl = nullptr;
+  ValueDecl *MainDecl = nullptr;
 
   /// The source location of the main type.
   SourceLoc MainDeclDiagLoc;
@@ -479,7 +479,7 @@ class SourceFile final : public FileUnit {
     llvm_unreachable("bad SourceFileKind");
   }
 
-  Decl *getMainDecl() const override { return MainDecl; }
+  ValueDecl *getMainDecl() const override { return MainDecl; }
   SourceLoc getMainDeclDiagLoc() const {
     assert(hasMainDecl());
     return MainDeclDiagLoc;
@@ -493,7 +493,7 @@ class SourceFile final : public FileUnit {
   /// one.
   ///
   /// Should only be called during type-checking.
-  bool registerMainDecl(Decl *mainDecl, SourceLoc diagLoc);
+  bool registerMainDecl(ValueDecl *mainDecl, SourceLoc diagLoc);
 
   /// True if this source file has an application entry point.
   ///

include/swift/AST/TypeAlignments.h

@@ -38,6 +38,7 @@ namespace swift {
   class DifferentiableAttr;
   class Expr;
   class ExtensionDecl;
+  class FileUnit;
   class GenericEnvironment;
   class GenericParamList;
   class GenericTypeParamDecl;
@@ -123,6 +124,7 @@ LLVM_DECLARE_TYPE_ALIGNMENT(swift::BraceStmt, swift::StmtAlignInBits)
 
 LLVM_DECLARE_TYPE_ALIGNMENT(swift::ASTContext, swift::ASTContextAlignInBits);
 LLVM_DECLARE_TYPE_ALIGNMENT(swift::DeclContext, swift::DeclContextAlignInBits)
+LLVM_DECLARE_TYPE_ALIGNMENT(swift::FileUnit, swift::DeclContextAlignInBits)
 LLVM_DECLARE_TYPE_ALIGNMENT(swift::DifferentiableAttr, swift::PointerAlignInBits)
 LLVM_DECLARE_TYPE_ALIGNMENT(swift::Expr, swift::ExprAlignInBits)
 LLVM_DECLARE_TYPE_ALIGNMENT(swift::CaptureListExpr, swift::ExprAlignInBits)

include/swift/SIL/SILDeclRef.h

@@ -42,6 +42,7 @@ namespace swift {
   class AutoClosureExpr;
   class ASTContext;
   class ClassDecl;
+  class FileUnit;
   class SILFunctionType;
   enum IsSerialized_t : unsigned char;
   enum class SubclassScope : unsigned char;
@@ -83,7 +84,14 @@ enum ForDefinition_t : bool {
 /// declaration, such as uncurry levels of a function, the allocating and
 /// initializing entry points of a constructor, etc.
 struct SILDeclRef {
-  using Loc = llvm::PointerUnion<ValueDecl *, AbstractClosureExpr *>;
+  /// The type of AST node location being stored.
+  enum LocKind {
+    Decl,
+    Closure,
+    File
+  };
+  using Loc = llvm::PointerUnion<ValueDecl *, AbstractClosureExpr *,
+                                 FileUnit *>;
 
   /// Represents the "kind" of the SILDeclRef. For some Swift decls there
   /// are multiple SIL entry points, and the kind is used to distinguish them.
@@ -144,9 +152,13 @@ struct SILDeclRef {
     /// References the function used to initialize a property wrapper storage
     /// instance from a projected value.
     PropertyWrapperInitFromProjectedValue,
+
+    /// The main entry-point function. This may reference a SourceFile for a
+    /// top-level main, or a decl for e.g an @main decl.
+    EntryPoint,
   };
 
-  /// The ValueDecl or AbstractClosureExpr represented by this SILDeclRef.
+  /// The AST node represented by this SILDeclRef.
   Loc loc;
   /// The Kind of this SILDeclRef.
   Kind kind : 4;
@@ -159,6 +171,17 @@ struct SILDeclRef {
                const GenericSignatureImpl *>
       pointer;
 
+  /// Returns the type of AST node location being stored by the SILDeclRef.
+  LocKind getLocKind() const {
+    if (loc.is<ValueDecl *>())
+      return LocKind::Decl;
+    if (loc.is<AbstractClosureExpr *>())
+      return LocKind::Closure;
+    if (loc.is<FileUnit *>())
+      return LocKind::File;
+    llvm_unreachable("Unhandled location kind!");
+  }
+
   /// The derivative function identifier.
   AutoDiffDerivativeFunctionIdentifier * getDerivativeFunctionIdentifier() const {
     if (!pointer.is<AutoDiffDerivativeFunctionIdentifier *>())
@@ -201,6 +224,12 @@ struct SILDeclRef {
   /// Produce a SIL constant for a default argument generator.
   static SILDeclRef getDefaultArgGenerator(Loc loc, unsigned defaultArgIndex);
 
+  /// Produces a SILDeclRef for a synthetic main entry-point such as @main.
+  static SILDeclRef getMainDeclEntryPoint(ValueDecl *decl);
+
+  /// Produces a SILDeclRef for the entry-point of a main FileUnit.
+  static SILDeclRef getMainFileEntryPoint(FileUnit *file);
+
   bool isNull() const { return loc.isNull(); }
   explicit operator bool() const { return !isNull(); }
 
@@ -217,7 +246,13 @@ struct SILDeclRef {
   AutoClosureExpr *getAutoClosureExpr() const;
   FuncDecl *getFuncDecl() const;
   AbstractFunctionDecl *getAbstractFunctionDecl() const;
-
+  FileUnit *getFileUnit() const {
+    return loc.get<FileUnit *>();
+  }
+
+  /// Retrieves the ASTContext from the underlying AST node being stored.
+  ASTContext &getASTContext() const;
+
   llvm::Optional<AnyFunctionRef> getAnyFunctionRef() const;
 
   SILLocation getAsRegularLocation() const;

include/swift/Serialization/SerializedModuleLoader.h

@@ -432,7 +432,7 @@ class SerializedASTFile final : public LoadedFile {
 
   virtual StringRef getModuleDefiningPath() const override;
 
-  Decl *getMainDecl() const override;
+  ValueDecl *getMainDecl() const override;
 
   bool hasEntryPoint() const override;
 

lib/AST/Module.cpp

@@ -1418,7 +1418,7 @@ bool ModuleDecl::isBuiltinModule() const {
   return this == getASTContext().TheBuiltinModule;
 }
 
-bool SourceFile::registerMainDecl(Decl *mainDecl, SourceLoc diagLoc) {
+bool SourceFile::registerMainDecl(ValueDecl *mainDecl, SourceLoc diagLoc) {
   assert(mainDecl);
   if (mainDecl == MainDecl)
     return false;

lib/IRGen/GenObjC.cpp

@@ -473,6 +473,7 @@ namespace {
       case SILDeclRef::Kind::GlobalAccessor:
       case SILDeclRef::Kind::PropertyWrapperBackingInitializer:
       case SILDeclRef::Kind::PropertyWrapperInitFromProjectedValue:
+      case SILDeclRef::Kind::EntryPoint:
         llvm_unreachable("Method does not have a selector");
 
       case SILDeclRef::Kind::Destroyer:

lib/SIL/IR/SILDeclRef.cpp

@@ -16,6 +16,7 @@
 #include "swift/AST/AnyFunctionRef.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/ParameterList.h"
+#include "swift/AST/SourceFile.h"
 #include "swift/ClangImporter/ClangImporter.h"
 #include "swift/ClangImporter/ClangModule.h"
 #include "swift/SIL/SILLinkage.h"
@@ -170,14 +171,29 @@ SILDeclRef::SILDeclRef(SILDeclRef::Loc baseLoc,
 }
 
 Optional<AnyFunctionRef> SILDeclRef::getAnyFunctionRef() const {
-  if (auto vd = loc.dyn_cast<ValueDecl*>()) {
-    if (auto afd = dyn_cast<AbstractFunctionDecl>(vd)) {
+  switch (getLocKind()) {
+  case LocKind::Decl:
+    if (auto *afd = getAbstractFunctionDecl())
       return AnyFunctionRef(afd);
-    } else {
-      return None;
-    }
+    return None;
+  case LocKind::Closure:
+    return AnyFunctionRef(getAbstractClosureExpr());
+  case LocKind::File:
+    return None;
+  }
+  llvm_unreachable("Unhandled case in switch");
+}
+
+ASTContext &SILDeclRef::getASTContext() const {
+  switch (getLocKind()) {
+  case LocKind::Decl:
+    return getDecl()->getASTContext();
+  case LocKind::Closure:
+    return getAbstractClosureExpr()->getASTContext();
+  case LocKind::File:
+    return getFileUnit()->getASTContext();
   }
-  return AnyFunctionRef(loc.get<AbstractClosureExpr*>());
+  llvm_unreachable("Unhandled case in switch");
 }
 
 bool SILDeclRef::isThunk() const {
@@ -226,9 +242,16 @@ bool SILDeclRef::isClangGenerated(ClangNode node) {
 }
 
 bool SILDeclRef::isImplicit() const {
-  if (hasDecl())
+  switch (getLocKind()) {
+  case LocKind::Decl:
     return getDecl()->isImplicit();
-  return getAbstractClosureExpr()->isImplicit();
+  case LocKind::Closure:
+    return getAbstractClosureExpr()->isImplicit();
+  case LocKind::File:
+    // Files are currently never considered implicit.
+    return false;
+  }
+  llvm_unreachable("Unhandled case in switch");
 }
 
 SILLinkage SILDeclRef::getLinkage(ForDefinition_t forDefinition) const {
@@ -242,6 +265,10 @@ SILLinkage SILDeclRef::getLinkage(ForDefinition_t forDefinition) const {
     return isSerialized() ? SILLinkage::Shared : SILLinkage::Private;
   }
 
+  // The main entry-point is public.
+  if (kind == Kind::EntryPoint)
+    return SILLinkage::Public;
+
   // Add External to the linkage (e.g. Public -> PublicExternal) if this is a
   // declaration not a definition.
   auto maybeAddExternal = [&](SILLinkage linkage) {
@@ -407,6 +434,23 @@ SILDeclRef SILDeclRef::getDefaultArgGenerator(Loc loc,
   return result;
 }
 
+SILDeclRef SILDeclRef::getMainDeclEntryPoint(ValueDecl *decl) {
+  auto *file = cast<FileUnit>(decl->getDeclContext()->getModuleScopeContext());
+  assert(file->getMainDecl() == decl);
+  SILDeclRef result;
+  result.loc = decl;
+  result.kind = Kind::EntryPoint;
+  return result;
+}
+
+SILDeclRef SILDeclRef::getMainFileEntryPoint(FileUnit *file) {
+  assert(file->hasEntryPoint() && !file->getMainDecl());
+  SILDeclRef result;
+  result.loc = file;
+  result.kind = Kind::EntryPoint;
+  return result;
+}
+
 bool SILDeclRef::hasClosureExpr() const {
   return loc.is<AbstractClosureExpr *>()
     && isa<ClosureExpr>(getAbstractClosureExpr());
@@ -490,6 +534,9 @@ IsSerialized_t SILDeclRef::isSerialized() const {
     return IsNotSerialized;
   }
 
+  if (kind == Kind::EntryPoint)
+    return IsNotSerialized;
+
   if (isIVarInitializerOrDestroyer())
     return IsNotSerialized;
 
@@ -681,17 +728,23 @@ bool SILDeclRef::isNativeToForeignThunk() const {
   if (!isForeign)
     return false;
 
-  // We can have native-to-foreign thunks over closures.
-  if (!hasDecl())
-    return true;
+  switch (getLocKind()) {
+  case LocKind::Decl:
+    // A decl with a clang node doesn't have a native entry-point to forward
+    // onto.
+    if (getDecl()->hasClangNode())
+      return false;
 
-  // A decl with a clang node doesn't have a native entry-point to forward onto.
-  if (getDecl()->hasClangNode())
+    // Only certain kinds of SILDeclRef can expose native-to-foreign thunks.
+    return kind == Kind::Func || kind == Kind::Initializer ||
+           kind == Kind::Deallocator;
+  case LocKind::Closure:
+    // We can have native-to-foreign thunks over closures.
+    return true;
+  case LocKind::File:
     return false;
-
-  // Only certain kinds of SILDeclRef can expose native-to-foreign thunks.
-  return kind == Kind::Func || kind == Kind::Initializer ||
-         kind == Kind::Deallocator;
+  }
+  llvm_unreachable("Unhandled case in switch");
 }
 
 /// Use the Clang importer to mangle a Clang declaration.
@@ -784,8 +837,8 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
 
   switch (kind) {
   case SILDeclRef::Kind::Func:
-    if (!hasDecl())
-      return mangler.mangleClosureEntity(getAbstractClosureExpr(), SKind);
+    if (auto *ACE = getAbstractClosureExpr())
+      return mangler.mangleClosureEntity(ACE, SKind);
 
     // As a special case, functions can have manually mangled names.
     // Use the SILGen name only for the original non-thunked, non-curried entry
@@ -856,6 +909,10 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
   case SILDeclRef::Kind::PropertyWrapperInitFromProjectedValue:
     return mangler.mangleInitFromProjectedValueEntity(cast<VarDecl>(getDecl()),
                                                       SKind);
+
+  case SILDeclRef::Kind::EntryPoint: {
+    return getASTContext().getEntryPointFunctionName();
+  }
   }
 
   llvm_unreachable("bad entity kind!");
@@ -1059,9 +1116,15 @@ SILDeclRef SILDeclRef::getOverriddenVTableEntry() const {
 }
 
 SILLocation SILDeclRef::getAsRegularLocation() const {
-  if (hasDecl())
+  switch (getLocKind()) {
+  case LocKind::Decl:
     return RegularLocation(getDecl());
-  return RegularLocation(getAbstractClosureExpr());
+  case LocKind::Closure:
+    return RegularLocation(getAbstractClosureExpr());
+  case LocKind::File:
+    return RegularLocation::getModuleLocation();
+  }
+  llvm_unreachable("Unhandled case in switch");
 }
 
 SubclassScope SILDeclRef::getSubclassScope() const {
@@ -1168,7 +1231,12 @@ SubclassScope SILDeclRef::getSubclassScope() const {
 }
 
 unsigned SILDeclRef::getParameterListCount() const {
-  if (!hasDecl() || kind == Kind::DefaultArgGenerator)
+  // Only decls can introduce currying.
+  if (!hasDecl())
+    return 1;
+
+  // Always uncurried even if the underlying function is curried.
+  if (kind == Kind::DefaultArgGenerator || kind == Kind::EntryPoint)
     return 1;
 
   auto *vd = getDecl();