[Driver][Index] Add driver support to specify an overriding output path to record in the index data

include/swift/Basic/FileTypes.def

@@ -79,6 +79,7 @@ TYPE("json-features",       JSONFeatures,              "features.json",   "")
 
 
 TYPE("index-data",          IndexData,                 "",                "")
+TYPE("index-unit-output-path", IndexUnitOutputPath,    "",                "")
 TYPE("yaml-opt-record",     YAMLOptRecord,             "opt.yaml",        "")
 TYPE("bitstream-opt-record",BitstreamOptRecord,        "opt.bitstream",   "")
 

include/swift/Driver/Job.h

@@ -157,7 +157,8 @@ class CommandOutput {
 
   /// Associate a new \p PrimaryOutputFile (of type \c getPrimaryOutputType())
   /// with the provided \p Input pair of Base and Primary inputs.
-  void addPrimaryOutput(CommandInputPair Input, StringRef PrimaryOutputFile);
+  void addPrimaryOutput(CommandInputPair Input, StringRef PrimaryOutputFile,
+                        StringRef IndexUnitOutputPath);
 
   /// Return true iff the set of additional output types in \c this is
   /// identical to the set of additional output types in \p other.
@@ -184,6 +185,13 @@ class CommandOutput {
   /// assumed at several call sites.
   SmallVector<StringRef, 16> getPrimaryOutputFilenames() const;
 
+
+  /// Returns the output file path to record in the index data for each input.
+  /// The return value will contain one \c StringRef per primary input if any
+  /// input had an output filename for the index data that was different to its
+  /// primary output filename, and be empty otherwise.
+  SmallVector<StringRef, 16> getIndexUnitOutputFilenames() const;
+
   /// Assuming (and asserting) that there are one or more input pairs, associate
   /// an additional output named \p OutputFilename of type \p type with the
   /// first primary input. If the provided \p type is the primary output type,

include/swift/Driver/Util.h

@@ -49,6 +49,7 @@ namespace driver {
       SourceInputActions,
       InputJobsAndSourceInputActions,
       Output,
+      IndexUnitOutputPaths,
       /// Batch mode frontend invocations may have so many supplementary
       /// outputs that they don't comfortably fit as command-line arguments.
       /// In that case, add a FilelistInfo to record the path to the file.

lib/Basic/FileTypes.cpp

@@ -105,6 +105,7 @@ bool file_types::isTextual(ID Id) {
   case file_types::TY_Remapping:
   case file_types::TY_IndexData:
   case file_types::TY_BitstreamOptRecord:
+  case file_types::TY_IndexUnitOutputPath:
     return false;
   case file_types::TY_INVALID:
     llvm_unreachable("Invalid type ID.");
@@ -155,6 +156,7 @@ bool file_types::isAfterLLVM(ID Id) {
   case file_types::TY_PrivateSwiftModuleInterfaceFile:
   case file_types::TY_JSONDependencies:
   case file_types::TY_JSONFeatures:
+  case file_types::TY_IndexUnitOutputPath:
     return false;
   case file_types::TY_INVALID:
     llvm_unreachable("Invalid type ID.");
@@ -205,6 +207,7 @@ bool file_types::isPartOfSwiftCompilation(ID Id) {
   case file_types::TY_BitstreamOptRecord:
   case file_types::TY_JSONDependencies:
   case file_types::TY_JSONFeatures:
+  case file_types::TY_IndexUnitOutputPath:
     return false;
   case file_types::TY_INVALID:
     llvm_unreachable("Invalid type ID.");

lib/Driver/Compilation.cpp

@@ -1699,6 +1699,12 @@ static void writeOutputToFilelist(llvm::raw_fd_ostream &out, const Job *job,
   for (auto &output : outputInfo.getPrimaryOutputFilenames())
     out << output << "\n";
 }
+static void writeIndexUnitOutputPathsToFilelist(llvm::raw_fd_ostream &out,
+                                                const Job *job) {
+  const CommandOutput &outputInfo = job->getOutput();
+  for (auto &output : outputInfo.getIndexUnitOutputFilenames())
+    out << output << "\n";
+}
 static void writeSupplementarOutputToFilelist(llvm::raw_fd_ostream &out,
                                               const Job *job) {
   job->getOutput().writeOutputFileMap(out);
@@ -1732,13 +1738,15 @@ static bool writeFilelistIfNecessary(const Job *job, const ArgList &args,
       writeInputJobsToFilelist(out, job, filelistInfo.type);
       writeSourceInputActionsToFilelist(out, job, args);
       break;
-    case FilelistInfo::WhichFiles::Output: {
+    case FilelistInfo::WhichFiles::Output:
       writeOutputToFilelist(out, job, filelistInfo.type);
       break;
-      }
-      case FilelistInfo::WhichFiles::SupplementaryOutput:
-        writeSupplementarOutputToFilelist(out, job);
-        break;
+    case FilelistInfo::WhichFiles::IndexUnitOutputPaths:
+      writeIndexUnitOutputPathsToFilelist(out, job);
+      break;
+    case FilelistInfo::WhichFiles::SupplementaryOutput:
+      writeSupplementarOutputToFilelist(out, job);
+      break;
     }
   }
   return ok;

lib/Driver/Driver.cpp

@@ -2042,6 +2042,7 @@ void Driver::buildActions(SmallVectorImpl<const Action *> &TopLevelActions,
       case file_types::TY_RawSIB:
       case file_types::TY_RawSIL:
       case file_types::TY_Nothing:
+      case file_types::TY_IndexUnitOutputPath:
       case file_types::TY_INVALID:
         llvm_unreachable("these types should never be inferred");
       }
@@ -2507,6 +2508,28 @@ static StringRef baseNameForImage(const JobAction *JA, const OutputInfo &OI,
   return Buffer.str();
 }
 
+static StringRef getIndexUnitOutputFilename(Compilation &C,
+                                            const JobAction *JA,
+                                            const TypeToPathMap *OutputMap,
+                                            bool AtTopLevel) {
+  if (JA->getType() == file_types::TY_Nothing)
+    return {};
+
+  if (OutputMap) {
+    auto iter = OutputMap->find(file_types::TY_IndexUnitOutputPath);
+    if (iter != OutputMap->end())
+      return iter->second;
+  }
+
+  if (AtTopLevel) {
+    const llvm::opt::DerivedArgList &Args = C.getArgs();
+    if (Arg *FinalOutput = Args.getLastArg(options::OPT_index_unit_output_path))
+      return FinalOutput->getValue();
+  }
+
+  return {};
+}
+
 static StringRef getOutputFilename(Compilation &C,
                                    const JobAction *JA,
                                    const TypeToPathMap *OutputMap,
@@ -3021,7 +3044,7 @@ void Driver::computeMainOutput(
     SmallVectorImpl<const Job *> &InputJobs, const TypeToPathMap *OutputMap,
     StringRef workingDirectory, StringRef BaseInput, StringRef PrimaryInput,
     llvm::SmallString<128> &Buf, CommandOutput *Output) const {
-  StringRef OutputFile;
+  StringRef OutputFile, IndexUnitOutputFile;
   if (C.getOutputInfo().isMultiThreading() && isa<CompileJobAction>(JA) &&
       file_types::isAfterLLVM(JA->getType())) {
     // Multi-threaded compilation: A single frontend command produces multiple
@@ -3042,8 +3065,10 @@ void Driver::computeMainOutput(
 
       OutputFile = getOutputFilename(C, JA, OMForInput, workingDirectory,
                                      AtTopLevel, Base, Primary, Buf);
+      IndexUnitOutputFile = getIndexUnitOutputFilename(C, JA, OMForInput,
+                                                       AtTopLevel);
       Output->addPrimaryOutput(CommandInputPair(Base, Primary),
-                               OutputFile);
+                               OutputFile, IndexUnitOutputFile);
     };
     // Add an output file for each input action.
     for (const Action *A : InputActions) {
@@ -3063,8 +3088,10 @@ void Driver::computeMainOutput(
     // The common case: there is a single output file.
     OutputFile = getOutputFilename(C, JA, OutputMap, workingDirectory,
                                    AtTopLevel, BaseInput, PrimaryInput, Buf);
+    IndexUnitOutputFile = getIndexUnitOutputFilename(C, JA, OutputMap,
+                                                     AtTopLevel);
     Output->addPrimaryOutput(CommandInputPair(BaseInput, PrimaryInput),
-                             OutputFile);
+                             OutputFile, IndexUnitOutputFile);
   }
 }
 

lib/Driver/Job.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/Basic/Range.h"
 #include "swift/Basic/STLExtras.h"
 #include "swift/Driver/Job.h"
 #include "swift/Driver/PrettyStackTrace.h"
@@ -145,7 +146,8 @@ file_types::ID CommandOutput::getPrimaryOutputType() const {
 }
 
 void CommandOutput::addPrimaryOutput(CommandInputPair Input,
-                                     StringRef PrimaryOutputFile) {
+                                     StringRef PrimaryOutputFile,
+                                     StringRef IndexUnitOutputPath) {
   PrettyStackTraceDriverCommandOutputAddition CrashInfo(
     "primary", this, Input.Primary, PrimaryOutputType, PrimaryOutputFile);
   if (PrimaryOutputType == file_types::TY_Nothing) {
@@ -165,6 +167,13 @@ void CommandOutput::addPrimaryOutput(CommandInputPair Input,
   assert(!PrimaryOutputFile.empty());
   assert(AdditionalOutputTypes.count(PrimaryOutputType) == 0);
   ensureEntry(Input.Primary, PrimaryOutputType, PrimaryOutputFile, false);
+
+  // If there is an overriding output path to record in the index store instead
+  // of the primary output path, add it.
+  if (!IndexUnitOutputPath.empty()) {
+    ensureEntry(Input.Primary, file_types::TY_IndexUnitOutputPath,
+                IndexUnitOutputPath, false);
+  }
 }
 
 StringRef CommandOutput::getPrimaryOutputFilename() const {
@@ -188,6 +197,34 @@ SmallVector<StringRef, 16> CommandOutput::getPrimaryOutputFilenames() const {
   return V;
 }
 
+SmallVector<StringRef, 16>
+CommandOutput::getIndexUnitOutputFilenames() const {
+  SmallVector<StringRef, 16> V;
+  size_t NonEmpty = 0;
+  for (auto const &I: Inputs) {
+    auto Out = getOutputForInputAndType(I.Primary,
+                                        file_types::TY_IndexUnitOutputPath);
+    V.push_back(Out);
+    if (!Out.empty())
+      ++NonEmpty;
+  }
+
+  if (!NonEmpty || NonEmpty == V.size()) {
+    if (!NonEmpty)
+      V.clear();
+    return V;
+  }
+
+  auto PrimaryOutputs = getPrimaryOutputFilenames();
+  assert(PrimaryOutputs.size() == V.size());
+
+  for (auto I: indices(V)) {
+    if (V[I].empty())
+      V[I] = PrimaryOutputs[I];
+  }
+  return V;
+}
+
 void CommandOutput::setAdditionalOutputForType(file_types::ID Type,
                                                StringRef OutputFilename) {
   PrettyStackTraceDriverCommandOutputAddition CrashInfo(

lib/Driver/ToolChains.cpp

@@ -498,17 +498,32 @@ ToolChain::constructInvocation(const CompileJobAction &job,
 
   // Add the output file argument if necessary.
   if (context.Output.getPrimaryOutputType() != file_types::TY_Nothing) {
+    auto IndexUnitOutputs = context.Output.getIndexUnitOutputFilenames();
+
     if (context.shouldUseMainOutputFileListInFrontendInvocation()) {
       Arguments.push_back("-output-filelist");
       Arguments.push_back(context.getTemporaryFilePath("outputs", ""));
       II.FilelistInfos.push_back({Arguments.back(),
                                   context.Output.getPrimaryOutputType(),
                                   FilelistInfo::WhichFiles::Output});
+      if (!IndexUnitOutputs.empty()) {
+        Arguments.push_back("-index-unit-output-path-filelist");
+        Arguments.push_back(context.getTemporaryFilePath("index-unit-outputs",
+                                                         ""));
+        II.FilelistInfos.push_back({
+          Arguments.back(), file_types::TY_IndexUnitOutputPath,
+          FilelistInfo::WhichFiles::IndexUnitOutputPaths});
+      }
     } else {
       for (auto FileName : context.Output.getPrimaryOutputFilenames()) {
         Arguments.push_back("-o");
         Arguments.push_back(context.Args.MakeArgString(FileName));
       }
+
+      for (auto FileName : IndexUnitOutputs) {
+        Arguments.push_back("-index-unit-output-path");
+        Arguments.push_back(context.Args.MakeArgString(FileName));
+      }
     }
   }
 
@@ -639,6 +654,7 @@ const char *ToolChain::JobContext::computeFrontendModeForCompile() const {
   case file_types::TY_SwiftSourceInfoFile:
   case file_types::TY_SwiftCrossImportDir:
   case file_types::TY_SwiftOverlayFile:
+  case file_types::TY_IndexUnitOutputPath:
     llvm_unreachable("Output type can never be primary output.");
   case file_types::TY_INVALID:
     llvm_unreachable("Invalid type ID");
@@ -896,6 +912,7 @@ ToolChain::constructInvocation(const BackendJobAction &job,
     case file_types::TY_SwiftSourceInfoFile:
     case file_types::TY_SwiftCrossImportDir:
     case file_types::TY_SwiftOverlayFile:
+    case file_types::TY_IndexUnitOutputPath:
       llvm_unreachable("Output type can never be primary output.");
     case file_types::TY_INVALID:
       llvm_unreachable("Invalid type ID");

lib/Option/features.json

@@ -5,6 +5,9 @@
     },
     {
       "name": "experimental-allow-module-with-compiler-errors"
+    },
+    {
+      "name": "index-unit-output-path"
     }
   ]
 }