[NFC] Refactor Main File Computation

Author: CodaFi

include/swift/Frontend/Frontend.h

@@ -430,9 +430,6 @@ class CompilerInstance {
   /// \c mutable as it is consumed by \c loadPartialModulesAndImplicitImports.
   mutable std::vector<ModuleBuffers> PartialModules;
 
-  enum : unsigned { NO_SUCH_BUFFER = ~0U };
-  unsigned MainBufferID = NO_SUCH_BUFFER;
-
   /// Identifies the set of input buffers in the SourceManager that are
   /// considered primaries.
   llvm::SetVector<unsigned> PrimaryBufferIDs;
@@ -602,12 +599,14 @@ class CompilerInstance {
   /// Creates a new source file for the main module.
   SourceFile *createSourceFileForMainModule(ModuleDecl *mod,
                                             SourceFileKind FileKind,
-                                            Optional<unsigned> BufferID) const;
+                                            Optional<unsigned> BufferID,
+                                            bool isMainBuffer = false) const;
 
   /// Creates all the files to be added to the main module, appending them to
   /// \p files. If a loading error occurs, returns \c true.
   bool createFilesForMainModule(ModuleDecl *mod,
                                 SmallVectorImpl<FileUnit *> &files) const;
+  SourceFile *computeMainSourceFileForModule(ModuleDecl *mod) const;
 
 public:
   void freeASTContext();

lib/Frontend/Frontend.cpp

@@ -578,21 +578,6 @@ SourceFile *CompilerInstance::getCodeCompletionFile() const {
   return evaluateOrDefault(eval, CodeCompletionFileRequest{mod}, nullptr);
 }
 
-static bool shouldTreatSingleInputAsMain(InputFileKind inputKind) {
-  switch (inputKind) {
-  case InputFileKind::Swift:
-  case InputFileKind::SwiftModuleInterface:
-  case InputFileKind::SIL:
-    return true;
-  case InputFileKind::SwiftLibrary:
-  case InputFileKind::LLVM:
-  case InputFileKind::ObjCHeader:
-  case InputFileKind::None:
-    return false;
-  }
-  llvm_unreachable("unhandled input kind");
-}
-
 bool CompilerInstance::setUpInputs() {
   // Adds to InputSourceCodeBufferIDs, so may need to happen before the
   // per-input setup.
@@ -619,15 +604,6 @@ bool CompilerInstance::setUpInputs() {
     recordPrimaryInputBuffer(*codeCompletionBufferID);
   }
 
-  if (MainBufferID == NO_SUCH_BUFFER &&
-      InputSourceCodeBufferIDs.size() == 1 &&
-      shouldTreatSingleInputAsMain(Invocation.getInputKind())) {
-    MainBufferID = InputSourceCodeBufferIDs.front();
-  }
-
-  return false;
-}
-
   return false;
 }
 
@@ -772,6 +748,62 @@ ImplicitImportInfo CompilerInstance::getImplicitImportInfo() const {
   return imports;
 }
 
+static Optional<SourceFileKind>
+tryMatchInputModeToSourceFileKind(FrontendOptions::ParseInputMode mode) {
+  switch (mode) {
+  case FrontendOptions::ParseInputMode::SwiftLibrary:
+      // A Swift file in -parse-as-library mode is a library file.
+    return SourceFileKind::Library;
+  case FrontendOptions::ParseInputMode::SIL:
+      // A Swift file in -parse-sil mode is a SIL file.
+    return SourceFileKind::SIL;
+  case FrontendOptions::ParseInputMode::SwiftModuleInterface:
+    return SourceFileKind::Interface;
+  case FrontendOptions::ParseInputMode::Swift:
+    return SourceFileKind::Main;
+  }
+  llvm::outs() << (unsigned)mode;
+  llvm_unreachable("Unhandled input parsing mode!");
+}
+
+SourceFile *
+CompilerInstance::computeMainSourceFileForModule(ModuleDecl *mod) const {
+  // Swift libraries cannot have a 'main'.
+  const auto &FOpts = getInvocation().getFrontendOptions();
+  const auto &Inputs = FOpts.InputsAndOutputs.getAllInputs();
+  if (FOpts.InputMode == FrontendOptions::ParseInputMode::SwiftLibrary) {
+    return nullptr;
+  }
+
+  // Try to pull out a file called 'main.swift'.
+  auto MainInputIter =
+      std::find_if(Inputs.begin(), Inputs.end(), [](const InputFile &input) {
+        return input.getType() == file_types::TY_Swift &&
+               llvm::sys::path::filename(input.getFileName()) == "main.swift";
+      });
+
+  Optional<unsigned> MainBufferID = None;
+  if (MainInputIter != Inputs.end()) {
+    MainBufferID =
+        getSourceMgr().getIDForBufferIdentifier(MainInputIter->getFileName());
+  } else if (InputSourceCodeBufferIDs.size() == 1) {
+    // Barring that, just nominate a single Swift file as the main file.
+    MainBufferID.emplace(InputSourceCodeBufferIDs.front());
+  }
+
+  if (!MainBufferID.hasValue()) {
+    return nullptr;
+  }
+
+  auto SFK = tryMatchInputModeToSourceFileKind(FOpts.InputMode);
+  if (!SFK.hasValue()) {
+    return nullptr;
+  }
+
+  return createSourceFileForMainModule(mod, *SFK,
+                                       *MainBufferID, /*isMainBuffer*/true);
+}
+
 bool CompilerInstance::createFilesForMainModule(
     ModuleDecl *mod, SmallVectorImpl<FileUnit *> &files) const {
   // Try to pull out the main source file, if any. This ensures that it
@@ -994,14 +1026,14 @@ CompilerInstance::getSourceFileParsingOptions(bool forPrimary) const {
 
 SourceFile *CompilerInstance::createSourceFileForMainModule(
     ModuleDecl *mod, SourceFileKind fileKind,
-    Optional<unsigned> bufferID) const {
+    Optional<unsigned> bufferID, bool isMainBuffer) const {
   auto isPrimary = bufferID && isPrimaryInput(*bufferID);
   auto opts = getSourceFileParsingOptions(isPrimary);
 
   auto *inputFile = new (*Context)
       SourceFile(*mod, fileKind, bufferID, opts, isPrimary);
 
-  if (bufferID == MainBufferID)
+  if (isMainBuffer)
     inputFile->SyntaxParsingCache = Invocation.getMainFileSyntaxParsingCache();
 
   return inputFile;

lib/FrontendTool/FrontendTool.cpp

@@ -1978,7 +1978,7 @@ static bool performCompile(CompilerInstance &Instance,
   const FrontendOptions::ActionType Action = opts.RequestedAction;
 
   // To compile LLVM IR, just pass it off unmodified.
-  if (Invocation.getInputKind() == InputFileKind::LLVM)
+  if (opts.InputsAndOutputs.shouldTreatAsLLVM())
     return compileLLVMIR(Instance);
 
   // If we aren't in a parse-only context and expect an implicit stdlib import,
@@ -1996,10 +1996,12 @@ static bool performCompile(CompilerInstance &Instance,
     if (FrontendOptions::shouldActionOnlyParse(Action)) {
       // Parsing gets triggered lazily, but let's make sure we have the right
       // input kind.
-      auto kind = Invocation.getInputKind();
-      return kind == InputFileKind::Swift ||
-             kind == InputFileKind::SwiftLibrary ||
-             kind == InputFileKind::SwiftModuleInterface;
+      return llvm::all_of(
+          opts.InputsAndOutputs.getAllInputs(), [](const InputFile &IF) {
+            const auto kind = IF.getType();
+            return kind == file_types::TY_Swift ||
+                   kind == file_types::TY_SwiftModuleInterfaceFile;
+          });
     }
     return true;
   }() && "Only supports parsing .swift files");