[BatchMode] Sink a bunch of FrontendTool uses of getPrimarySourceFile()

Author: graydon

include/swift/Frontend/Frontend.h

@@ -437,6 +437,16 @@ class CompilerInstance {
     return PrimarySourceFiles;
   }
 
+  /// Gets the Primary Source File if one exists, otherwise the main
+  /// module. If multiple Primary Source Files exist, fails with an
+  /// assertion.
+  ModuleOrSourceFile getPrimarySourceFileOrMainModule() {
+    if (PrimarySourceFiles.empty())
+      return getMainModule();
+    else
+      return getPrimarySourceFile();
+  }
+
   /// Gets the SourceFile which is the primary input for this CompilerInstance.
   /// \returns the primary SourceFile, or nullptr if there is no primary input;
   /// if there are _multiple_ primary inputs, fails with an assertion.

lib/FrontendTool/FrontendTool.cpp

@@ -453,8 +453,9 @@ static void countStatsPostSema(UnifiedStatsReporter &Stats,
     C.NumReferencedMemberNames = R->getUsedMembers().size();
   }
 
-  if (auto *SF = Instance.getPrimarySourceFile()) {
-    countStatsOfSourceFile(Stats, Instance, SF);
+  if (!Instance.getPrimarySourceFiles().empty()) {
+    for (auto SF : Instance.getPrimarySourceFiles())
+      countStatsOfSourceFile(Stats, Instance, SF);
   } else if (auto *M = Instance.getMainModule()) {
     // No primary source file, but a main module; this is WMO-mode
     for (auto *F : M->getFiles()) {
@@ -630,8 +631,6 @@ static bool performCompile(CompilerInstance &Instance,
     return Context.hadError();
   }
 
-  SourceFile *PrimarySourceFile = Instance.getPrimarySourceFile();
-
   // We've been told to dump the AST (either after parsing or type-checking,
   // which is already differentiated in CompilerInstance::performSema()),
   // so dump or print the main source file and return.
@@ -642,7 +641,7 @@ static bool performCompile(CompilerInstance &Instance,
       Action == FrontendOptions::ActionType::DumpScopeMaps ||
       Action == FrontendOptions::ActionType::DumpTypeRefinementContexts ||
       Action == FrontendOptions::ActionType::DumpInterfaceHash) {
-    SourceFile *SF = PrimarySourceFile;
+    SourceFile *SF = Instance.getPrimarySourceFile();
     if (!SF) {
       SourceFileKind Kind = Invocation.getSourceFileKind();
       SF = &Instance.getMainModule()->getMainSourceFile(Kind);
@@ -717,9 +716,12 @@ static bool performCompile(CompilerInstance &Instance,
     (void)emitMakeDependencies(Context.Diags, *Instance.getDependencyTracker(),
                                opts);
 
-  if (shouldTrackReferences)
-    emitReferenceDependencies(Context.Diags, Instance.getPrimarySourceFile(),
-                              *Instance.getDependencyTracker(), opts);
+  if (shouldTrackReferences) {
+    for (auto *SF : Instance.getPrimarySourceFiles()) {
+      emitReferenceDependencies(Context.Diags, SF,
+                                *Instance.getDependencyTracker(), opts);
+    }
+  }
 
   if (!opts.LoadedModuleTracePath.empty())
     (void)emitLoadedModuleTrace(Context, *Instance.getDependencyTracker(),
@@ -730,7 +732,8 @@ static bool performCompile(CompilerInstance &Instance,
   if (Context.hadError()) {
     if (shouldIndex) {
       //  Emit the index store data even if there were compiler errors.
-      if (emitIndexData(PrimarySourceFile, Invocation, Instance))
+      if (emitIndexData(Instance.getPrimarySourceFile(),
+                        Invocation, Instance))
         return true;
     }
     return true;
@@ -749,7 +752,8 @@ static bool performCompile(CompilerInstance &Instance,
       return printAsObjC(opts.ObjCHeaderOutputPath, Instance.getMainModule(),
                          opts.ImplicitObjCHeaderPath, moduleIsPublic);
     if (shouldIndex) {
-      if (emitIndexData(PrimarySourceFile, Invocation, Instance))
+      if (emitIndexData(Instance.getPrimarySourceFile(),
+                        Invocation, Instance))
         return true;
     }
     return Context.hadError();
@@ -780,7 +784,7 @@ static bool performCompile(CompilerInstance &Instance,
       return SASTF && SASTF->isSIB();
     };
     if (opts.Inputs.hasPrimaryInputs()) {
-      FileUnit *PrimaryFile = PrimarySourceFile;
+      FileUnit *PrimaryFile = Instance.getPrimarySourceFile();
       if (!PrimaryFile) {
         for (FileUnit *fileUnit : Instance.getMainModule()->getFiles()) {
           if (auto SASTF = dyn_cast<SerializedASTFile>(fileUnit)) {
@@ -827,8 +831,7 @@ static bool performCompile(CompilerInstance &Instance,
     if (Invocation.getSILOptions().LinkMode == SILOptions::LinkAll)
       performSILLinking(SM.get(), true);
 
-    auto DC = PrimarySourceFile ? ModuleOrSourceFile(PrimarySourceFile) :
-                                  Instance.getMainModule();
+    auto DC = Instance.getPrimarySourceFileOrMainModule();
     if (!opts.ModuleOutputPath.empty()) {
       SerializationOptions serializationOpts;
       serializationOpts.OutputPath = opts.ModuleOutputPath.c_str();
@@ -884,8 +887,7 @@ static bool performCompile(CompilerInstance &Instance,
 
   auto SerializeSILModuleAction = [&]() {
     if (!opts.ModuleOutputPath.empty() || !opts.ModuleDocOutputPath.empty()) {
-      auto DC = PrimarySourceFile ? ModuleOrSourceFile(PrimarySourceFile)
-                                  : Instance.getMainModule();
+        auto DC = Instance.getPrimarySourceFileOrMainModule();
       if (!opts.ModuleOutputPath.empty()) {
         SerializationOptions serializationOpts;
         serializationOpts.OutputPath = opts.ModuleOutputPath.c_str();
@@ -956,8 +958,9 @@ static bool performCompile(CompilerInstance &Instance,
 
   // Get the main source file's private discriminator and attach it to
   // the compile unit's flags.
-  if (PrimarySourceFile) {
-    Identifier PD = PrimarySourceFile->getPrivateDiscriminator();
+  if (IRGenOpts.DebugInfoKind != IRGenDebugInfoKind::None &&
+      Instance.getPrimarySourceFile()) {
+    Identifier PD = Instance.getPrimarySourceFile()->getPrivateDiscriminator();
     if (!PD.empty())
       IRGenOpts.DWARFDebugFlags += (" -private-discriminator "+PD.str()).str();
   }
@@ -968,8 +971,7 @@ static bool performCompile(CompilerInstance &Instance,
   }
 
   if (Action == FrontendOptions::ActionType::EmitSIB) {
-    auto DC = PrimarySourceFile ? ModuleOrSourceFile(PrimarySourceFile) :
-                                  Instance.getMainModule();
+    auto DC = Instance.getPrimarySourceFileOrMainModule();
     if (!opts.ModuleOutputPath.empty()) {
       SerializationOptions serializationOpts;
       serializationOpts.OutputPath = opts.ModuleOutputPath.c_str();
@@ -988,7 +990,8 @@ static bool performCompile(CompilerInstance &Instance,
     if (Action == FrontendOptions::ActionType::MergeModules ||
         Action == FrontendOptions::ActionType::EmitModuleOnly) {
       if (shouldIndex) {
-        if (emitIndexData(PrimarySourceFile, Invocation, Instance))
+        if (emitIndexData(Instance.getPrimarySourceFile(),
+                          Invocation, Instance))
           return true;
       }
       return Context.hadError();
@@ -1019,7 +1022,8 @@ static bool performCompile(CompilerInstance &Instance,
   // TODO: remove once the frontend understands what action it should perform
   IRGenOpts.OutputKind = getOutputKind(Action);
   if (Action == FrontendOptions::ActionType::Immediate) {
-    assert(!PrimarySourceFile && "-i doesn't work in -primary-file mode");
+    assert(Instance.getPrimarySourceFiles().empty() &&
+           "-i doesn't work in -primary-file mode");
     IRGenOpts.UseJIT = true;
     IRGenOpts.DebugInfoKind = IRGenDebugInfoKind::Normal;
     const ProcessCmdLine &CmdLine = ProcessCmdLine(opts.ImmediateArgv.begin(),
@@ -1040,8 +1044,10 @@ static bool performCompile(CompilerInstance &Instance,
   auto &LLVMContext = getGlobalLLVMContext();
   std::unique_ptr<llvm::Module> IRModule;
   llvm::GlobalVariable *HashGlobal;
-  if (PrimarySourceFile) {
-    IRModule = performIRGeneration(IRGenOpts, *PrimarySourceFile, std::move(SM),
+  if (!Instance.getPrimarySourceFiles().empty()) {
+    IRModule = performIRGeneration(IRGenOpts,
+                                   *Instance.getPrimarySourceFile(),
+                                   std::move(SM),
                                    opts.getSingleOutputFilename(), LLVMContext,
                                    0, &HashGlobal);
   } else {
@@ -1054,7 +1060,7 @@ static bool performCompile(CompilerInstance &Instance,
   // Walk the AST for indexing after IR generation. Walking it before seems
   // to cause miscompilation issues.
   if (shouldIndex) {
-    if (emitIndexData(PrimarySourceFile, Invocation, Instance))
+    if (emitIndexData(Instance.getPrimarySourceFile(), Invocation, Instance))
       return true;
   }
 
@@ -1083,8 +1089,9 @@ static bool performCompile(CompilerInstance &Instance,
     const auto &SILOpts = Invocation.getSILOptions();
     const auto hasMultipleIGMs = SILOpts.hasMultipleIGMs();
     bool error;
-    if (PrimarySourceFile)
-      error = validateTBD(PrimarySourceFile, *IRModule, hasMultipleIGMs,
+    if (!Instance.getPrimarySourceFiles().empty())
+      error = validateTBD(Instance.getPrimarySourceFile(),
+                          *IRModule, hasMultipleIGMs,
                           allSymbols);
     else
       error = validateTBD(Instance.getMainModule(), *IRModule, hasMultipleIGMs,