FineModuleTracing: only collect fine-grained tracing for those source files that compiler promises to type check.

rdar://141065994

Author: nkcsgexi

lib/FrontendTool/Dependencies.h

@@ -23,6 +23,7 @@ class DependencyTracker;
 class FrontendOptions;
 class InputFile;
 class ModuleDecl;
+class CompilerInstance;
 
 /// Emit the names of the modules imported by \c mainModule.
 bool emitImportedModules(ModuleDecl *mainModule, const FrontendOptions &opts,
@@ -32,7 +33,7 @@ bool emitLoadedModuleTraceIfNeeded(ModuleDecl *mainModule,
                                    const FrontendOptions &opts,
                                    const InputFile &input);
 
-bool emitFineModuleTraceIfNeeded(ModuleDecl *mainModule,
+bool emitFineModuleTraceIfNeeded(CompilerInstance &Instance,
                                  const FrontendOptions &opts);
 
 } // end namespace swift

lib/FrontendTool/FrontendTool.cpp

@@ -1001,7 +1001,7 @@ static void performEndOfPipelineActions(CompilerInstance &Instance) {
         Instance.getMainModule(), Instance.getDependencyTracker(), opts);
 
     dumpAPIIfNeeded(Instance);
-    swift::emitFineModuleTraceIfNeeded(Instance.getMainModule(), opts);
+    swift::emitFineModuleTraceIfNeeded(Instance, opts);
   }
 
   // Contains the hadError checks internally, we still want to output the

lib/FrontendTool/LoadedModuleTrace.cpp

@@ -22,6 +22,7 @@
 #include "swift/Basic/FileTypes.h"
 #include "swift/Basic/JSONSerialization.h"
 #include "swift/Basic/SourceManager.h"
+#include "swift/Frontend/Frontend.h"
 #include "swift/Frontend/FrontendOptions.h"
 #include "swift/Frontend/ModuleInterfaceSupport.h"
 #include "swift/IDE/SourceEntityWalker.h"
@@ -853,7 +854,8 @@ class ObjcMethodReferenceCollector: public SourceEntityWalker {
   }
 };
 
-static void createFineModuleTraceFile(const InputFile &input, ModuleDecl *MD) {
+static void createFineModuleTraceFile(CompilerInstance &instance,
+                                      const InputFile &input) {
   StringRef tracePath = input.getFineModuleTracePath();
   if (tracePath.empty()) {
     // we basically rely on the passing down of module trace file path
@@ -862,27 +864,8 @@ static void createFineModuleTraceFile(const InputFile &input, ModuleDecl *MD) {
     // specifically.
     return;
   }
+  ModuleDecl *MD = instance.getMainModule();
   auto &ctx = MD->getASTContext();
-  std::vector<SourceFile*> filesToWalk;
-  for (auto *FU : MD->getFiles()) {
-    if (auto *SF = dyn_cast<SourceFile>(FU)) {
-      switch (SF->Kind) {
-      case swift::SourceFileKind::Library:
-      case swift::SourceFileKind::Main:
-      case swift::SourceFileKind::MacroExpansion:
-      case swift::SourceFileKind::DefaultArgument:
-        filesToWalk.push_back(SF);
-        LLVM_FALLTHROUGH;
-      case swift::SourceFileKind::SIL:
-      case swift::SourceFileKind::Interface:
-        continue;
-      }
-    }
-  }
-  // No source files to walk, abort.
-  if (filesToWalk.empty()) {
-    return;
-  }
   // Write output via atomic append.
   llvm::vfs::OutputConfig config;
   config.setAppend().setAtomicWrite();
@@ -893,10 +876,11 @@ static void createFineModuleTraceFile(const InputFile &input, ModuleDecl *MD) {
     return;
   }
   ObjcMethodReferenceCollector collector(MD);
-  for (auto *SF: filesToWalk) {
-    collector.setFileBeforeVisiting(SF);
-    collector.walk(*SF);
-  }
+  instance.forEachFileToTypeCheck([&](SourceFile& SF) {
+    collector.setFileBeforeVisiting(&SF);
+    collector.walk(SF);
+    return false;
+  });
 
   // print this json line.
   std::string stringBuffer;
@@ -915,14 +899,15 @@ static void createFineModuleTraceFile(const InputFile &input, ModuleDecl *MD) {
   }
 }
 
-bool swift::emitFineModuleTraceIfNeeded(ModuleDecl *mainModule,
+bool swift::emitFineModuleTraceIfNeeded(CompilerInstance &Instance,
                                         const FrontendOptions &opts) {
+  ModuleDecl *mainModule = Instance.getMainModule();
   ASTContext &ctxt = mainModule->getASTContext();
   assert(!ctxt.hadError() &&
          "We should've already exited earlier if there was an error.");
 
   opts.InputsAndOutputs.forEachInput([&](const InputFile &input) {
-    createFineModuleTraceFile(input, mainModule);
+    createFineModuleTraceFile(Instance, input);
     return true;
   });
   return false;