Merge pull request #61279 from apple/rdar100130950

[IRGen] Emit bit for @main into entry section.

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()) {

lib/IRGen/GenDecl.cpp

@@ -2092,6 +2092,14 @@ void IRGenerator::emitEntryPointInfo() {
   auto entrypointInfo = builder.beginStruct();
   entrypointInfo.addCompactFunctionReference(
       IGM.getAddrOfSILFunction(entrypoint, NotForDefinition));
+  uint32_t flags = 0;
+  enum EntryPointFlags : unsigned {
+    HasAtMainTypeFlag = 1 << 0,
+  };
+  if (auto *mainTypeDecl = SIL.getSwiftModule()->getMainTypeDecl()) {
+    flags |= HasAtMainTypeFlag;
+  }
+  entrypointInfo.addInt(IGM.Int32Ty, flags);
   auto var = entrypointInfo.finishAndCreateGlobal(
       "\x01l_entry_point", Alignment(4),
       /*isConstant*/ true, llvm::GlobalValue::PrivateLinkage);

test/IRGen/Inputs/entry-point-section/main.swift

@@ -0,0 +1,2 @@
+print("hello world")
+

test/IRGen/entrypoint-section-run-main_swift.cpp

@@ -0,0 +1,82 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-clang %s -std=c++11 -isysroot %sdk -o %t/main
+// RUN: %target-codesign %t/main
+// RUN: %target-swift-frontend -c %S/Inputs/entry-point-section/main.swift -O -o %t/howdy.o -module-name Howdy
+// RUN: %target-ld %t/howdy.o -syslibroot %sdk -lSystem -dylib -o %t/libHowdy.dylib
+// RUN: %target-codesign %t/libHowdy.dylib
+// RUN: %target-run %t/main %t/libHowdy.dylib | %FileCheck %s
+
+// REQUIRES: VENDOR=apple
+// REQUIRES: executable_test
+// UNSUPPORTED: remote_run
+
+#include <dlfcn.h>
+#include <mach-o/dyld.h>
+#include <mach-o/getsect.h>
+#include <ptrauth.h>
+#include <stdio.h>
+#include <string.h>
+
+#if __POINTER_WIDTH__ == 64
+using mach_header_platform = mach_header_64;
+#else
+using mach_header_platform = mach_header;
+#endif
+
+int main(int argc, char *argv[]) {
+  if (argc != 2) {
+    printf("no argument!\n");
+    return 1;
+  }
+  auto *dylibPath = argv[1];
+  auto *handle = dlopen(dylibPath, RTLD_LAZY);
+  if (!handle) {
+    printf("no library!\n");
+    return 1;
+  }
+
+  using MainFunction = void(int, char *[]);
+  MainFunction *mainFunction = nullptr;
+  for (int index = 0; index < _dyld_image_count(); ++index) {
+    auto *imageName = _dyld_get_image_name(index);
+    if (strcmp(dylibPath, imageName)) {
+      printf("skipping %s\n", imageName);
+      continue;
+    }
+    auto *header = reinterpret_cast<const mach_header_platform *>(
+        _dyld_get_image_header(index));
+    size_t size;
+    auto *data = getsectiondata(header, "__TEXT", "__swift5_entry", &size);
+    int32_t offset = *reinterpret_cast<int32_t *>(data);
+    if (size >= 8) {
+      auto flags = *(reinterpret_cast<int32_t *>(data) + 1);
+      enum EntryPointFlags : unsigned {
+        HasAtMainTypeFlag = 1 << 0,
+      };
+      printf("flags: %d\n", flags);
+      bool isAtMainEntryPoint = flags & HasAtMainTypeFlag;
+      if (!isAtMainEntryPoint) {
+        printf("no @main entry point!\n"); // CHECK: no @main entry point!
+        continue;
+      }
+    }
+    mainFunction = reinterpret_cast<MainFunction *>(
+      ptrauth_sign_unauthenticated(
+          reinterpret_cast<void *>(
+              reinterpret_cast<long>(data) + offset
+          ),
+          ptrauth_key_function_pointer,
+          ptrauth_function_pointer_type_discriminator(MainFunction)
+      )
+    );
+
+
+    break;
+  }
+  if (!mainFunction) {
+    printf("no function!\n"); // CHECK: no function!
+    return 0;
+  }
+  mainFunction(argc, argv);
+  return 1;
+}

test/IRGen/entrypoint-section-run.cpp

@@ -47,6 +47,18 @@ int main(int argc, char *argv[]) {
     size_t size;
     auto *data = getsectiondata(header, "__TEXT", "__swift5_entry", &size);
     int32_t offset = *reinterpret_cast<int32_t *>(data);
+    if (size >= 8) {
+      auto flags = *(reinterpret_cast<int32_t *>(data) + 1);
+      enum EntryPointFlags : unsigned {
+        HasAtMainTypeFlag = 1 << 0,
+      };
+      printf("flags: %d\n", flags);
+      bool isAtMainEntryPoint = flags & HasAtMainTypeFlag;
+      if (!isAtMainEntryPoint) {
+        printf("no @main entry point!\n");
+        continue;
+      }
+    }
     mainFunction = reinterpret_cast<MainFunction *>(
       ptrauth_sign_unauthenticated(
           reinterpret_cast<void *>(
@@ -60,7 +72,7 @@ int main(int argc, char *argv[]) {
     break;
   }
   if (!mainFunction) {
-    printf("no function!");
+    printf("no function!\n");
     return 1;
   }
   mainFunction(argc, argv); // CHECK: howdy mundo

test/IRGen/unused.sil

@@ -52,11 +52,11 @@ bb0(%0 : $Int32, %1 : $UnsafeMutablePointer<Optional<UnsafeMutablePointer<Int8>>
   return %3 : $Int32                              // id: %4
 }
 
-// CHECK-macho: @"\01l_entry_point" = private constant { i32 } { i32 trunc (i64 sub (i64 ptrtoint (i32 (i32, i8**)* @main to i64), i64 ptrtoint ({ i32 }* @"\01l_entry_point" to i64)) to i32) }, section "__TEXT, __swift5_entry, regular, no_dead_strip", align 4
-// CHECK-elf: @"\01l_entry_point" = private constant { i32 } { i32 trunc (i64 sub (i64 ptrtoint (i32 (i32, i8**)* @main to i64), i64 ptrtoint ({ i32 }* @"\01l_entry_point" to i64)) to i32) }, section "swift5_entry", align 4
+// CHECK-macho: @"\01l_entry_point" = private constant { i32, i32 } { i32 trunc (i64 sub (i64 ptrtoint (i32 (i32, i8**)* @main to i64), i64 ptrtoint ({ i32, i32 }* @"\01l_entry_point" to i64)) to i32), i32 0 }, section "__TEXT, __swift5_entry, regular, no_dead_strip", align 4
+// CHECK-elf: @"\01l_entry_point" = private constant { i32, i32 } { i32 trunc (i64 sub (i64 ptrtoint (i32 (i32, i8**)* @main to i64), i64 ptrtoint ({ i32, i32 }* @"\01l_entry_point" to i64)) to i32), i32 0 }, section "swift5_entry", align 4
 
-// CHECK-macho: @llvm.used = appending global [4 x i8*] [i8* bitcast (void ()* @frieda to i8*), i8* bitcast (i32 (i32, i8**)* @main to i8*), i8* bitcast ({ i32 }* @"\01l_entry_point" to i8*), i8* bitcast (i16* @__swift_reflection_version to i8*)], section "llvm.metadata"
-// CHECK-elf: @llvm.compiler.used = appending global [5 x i8*] [i8* bitcast (void ()* @frieda to i8*), i8* bitcast (i32 (i32, i8**)* @main to i8*), i8* bitcast ({ i32 }* @"\01l_entry_point" to i8*), i8* bitcast (i16* @__swift_reflection_version to i8*), i8* getelementptr inbounds ([{{[0-9]+}} x i8], [{{[0-9]+}} x i8]* @_swift1_autolink_entries, i32 0, i32 0)], section "llvm.metadata"
+// CHECK-macho: @llvm.used = appending global [4 x i8*] [i8* bitcast (void ()* @frieda to i8*), i8* bitcast (i32 (i32, i8**)* @main to i8*), i8* bitcast ({ i32, i32 }* @"\01l_entry_point" to i8*), i8* bitcast (i16* @__swift_reflection_version to i8*)], section "llvm.metadata"
+// CHECK-elf: @llvm.compiler.used = appending global [5 x i8*] [i8* bitcast (void ()* @frieda to i8*), i8* bitcast (i32 (i32, i8**)* @main to i8*), i8* bitcast ({ i32, i32 }* @"\01l_entry_point" to i8*), i8* bitcast (i16* @__swift_reflection_version to i8*), i8* getelementptr inbounds ([{{[0-9]+}} x i8], [{{[0-9]+}} x i8]* @_swift1_autolink_entries, i32 0, i32 0)], section "llvm.metadata"
 
 // CHECK: define linkonce_odr hidden swiftcc void @qux()
 // CHECK: define hidden swiftcc void @fred()

test/lit.cfg

@@ -193,7 +193,7 @@ config.test_format = swift_test.SwiftTest(coverage_mode=config.coverage_mode,
 
 # suffixes: A list of file extensions to treat as test files.
 config.suffixes = ['.swift', '.ll', '.sil', '.gyb', '.m', '.c',
-                   '.swiftinterface', '.test-sh', '.test']
+                   '.swiftinterface', '.test-sh', '.test', '.cpp']
 
 # excludes: A list of directories to exclude from the testsuite. The 'Inputs'
 # subdirectories contain auxiliary inputs for various tests in their parent