[TBDGen] Use proper symbol source for main

Now that SILDeclRef can represent the main function,
tweak TBDGen to refer to it.

Author: hamishknight

lib/TBDGen/TBDGen.cpp

@@ -1152,6 +1152,29 @@ void TBDGenVisitor::addFirstFileSymbols() {
   }
 }
 
+void TBDGenVisitor::addMainIfNecessary(FileUnit *file) {
+  // HACK: 'main' is a special symbol that's always emitted in SILGen if
+  //       the file has an entry point. Since it doesn't show up in the
+  //       module until SILGen, we need to explicitly add it here.
+  //
+  // Make sure to only add the main symbol for the module that we're emitting
+  // TBD for, and not for any statically linked libraries.
+  if (!file->hasEntryPoint() || file->getParentModule() != SwiftModule)
+    return;
+
+  auto entryPointSymbol =
+      SwiftModule->getASTContext().getEntryPointFunctionName();
+
+  if (auto *decl = file->getMainDecl()) {
+    auto ref = SILDeclRef::getMainDeclEntryPoint(decl);
+    addSymbol(entryPointSymbol, SymbolSource::forSILDeclRef(ref));
+    return;
+  }
+
+  auto ref = SILDeclRef::getMainFileEntryPoint(file);
+  addSymbol(entryPointSymbol, SymbolSource::forSILDeclRef(ref));
+}
+
 void TBDGenVisitor::visit(Decl *D) {
   DeclStack.push_back(D);
   SWIFT_DEFER { DeclStack.pop_back(); };
@@ -1348,8 +1371,8 @@ class APIGenRecorder final : public APIRecorder {
     apigen::APIAvailability availability;
     if (source.kind == SymbolSource::Kind::SIL) {
       auto ref = source.getSILDeclRef();
-      if (auto *decl = ref.getDecl())
-        availability = getAvailability(decl);
+      if (ref.hasDecl())
+        availability = getAvailability(ref.getDecl());
     }
 
     api.addSymbol(symbol, moduleLoc, apigen::APILinkage::Exported,

lib/TBDGen/TBDGenVisitor.h

@@ -181,18 +181,9 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   TBDGenVisitor(const TBDGenDescriptor &desc, APIRecorder &recorder);
 
   ~TBDGenVisitor() { assert(DeclStack.empty()); }
-  void addMainIfNecessary(FileUnit *file) {
-    // HACK: 'main' is a special symbol that's always emitted in SILGen if
-    //       the file has an entry point. Since it doesn't show up in the
-    //       module until SILGen, we need to explicitly add it here.
-    //
-    // Make sure to only add the main symbol for the module that we're emitting
-    // TBD for, and not for any statically linked libraries.
-    // FIXME: We should have a SymbolSource for main.
-    if (file->hasEntryPoint() && file->getParentModule() == SwiftModule)
-      addSymbol(SwiftModule->getASTContext().getEntryPointFunctionName(),
-                SymbolSource::forUnknown());
-  }
+
+  /// Add the main symbol.
+  void addMainIfNecessary(FileUnit *file);
 
   /// Adds the global symbols associated with the first file.
   void addFirstFileSymbols();