[BatchMode] Add batch-mode support methods to ToolChain.

Author: graydon

include/swift/Driver/ToolChain.h

@@ -166,6 +166,30 @@ class ToolChain {
                                     std::unique_ptr<CommandOutput> output,
                                     const OutputInfo &OI) const;
 
+  /// Return true iff the input \c Job \p A is an acceptable candidate for
+  /// batching together into a BatchJob, via a call to \c
+  /// constructBatchJob. This is true when the \c Job is a built from a \c
+  /// CompileJobAction in a \c Compilation \p C running in \c
+  /// OutputInfo::Mode::StandardCompile output mode, with a single \c TY_Swift
+  /// \c InputAction.
+  bool jobIsBatchable(const Compilation &C, const Job *A) const;
+
+  /// Equivalence relation that holds iff the two input Jobs \p A and \p B are
+  /// acceptable candidates for combining together into a \c BatchJob, via a
+  /// call to \c constructBatchJob. This is true when each job independently
+  /// satisfies \c jobIsBatchable, and the two jobs have identical executables,
+  /// output types and environments (i.e. they are identical aside from their
+  /// inputs).
+  bool jobsAreBatchCombinable(const Compilation &C, const Job *A,
+                              const Job *B) const;
+
+  /// Construct a \c BatchJob that subsumes the work of a set of Jobs. Any pair
+  /// of elements in \p Jobs are assumed to satisfy the equivalence relation \c
+  /// jobsAreBatchCombinable, i.e. they should all be "the same" job in in all
+  /// ways other than their choices of inputs.
+  std::unique_ptr<Job> constructBatchJob(ArrayRef<const Job *> Jobs,
+                                         Compilation &C) const;
+
   /// Return the default language type to use for the given extension.
   /// If the extension is empty or is otherwise not recognized, return
   /// the invalid type \c TY_INVALID.

lib/Driver/ToolChain.cpp

@@ -21,6 +21,7 @@
 #include "swift/Driver/Compilation.h"
 #include "swift/Driver/Driver.h"
 #include "swift/Driver/Job.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
@@ -150,6 +151,170 @@ types::ID ToolChain::lookupTypeForExtension(StringRef Ext) const {
   return types::lookupTypeForExtension(Ext);
 }
 
+/// Return a _single_ TY_Swift InputAction, if one exists;
+/// if 0 or >1 such inputs exist, return nullptr.
+static const InputAction*
+findSingleSwiftInput(const CompileJobAction *CJA) {
+  auto Inputs = CJA->getInputs();
+  const InputAction *IA = nullptr;
+  for (auto const *I : Inputs) {
+    if (auto const *S = dyn_cast<InputAction>(I)) {
+      if (S->getType() == types::TY_Swift) {
+        if (IA == nullptr) {
+          IA = S;
+        } else {
+          // Already found one, two is too many.
+          return nullptr;
+        }
+      }
+    }
+  }
+  return IA;
+}
+
+static bool
+jobsHaveSameExecutableNames(const Job *A, const Job *B) {
+  // Jobs that get here (that are derived from CompileJobActions) should always
+  // have the same executable name -- it should always be SWIFT_EXECUTABLE_NAME
+  // -- but we check here just to be sure / fail gracefully in non-assert
+  // builds.
+  assert(strcmp(A->getExecutable(), B->getExecutable()) == 0);
+  if (strcmp(A->getExecutable(), B->getExecutable()) != 0) {
+    return false;
+  }
+  return true;
+}
+
+static bool
+jobsHaveSameOutputTypes(const Job *A, const Job *B) {
+  if (A->getOutput().getPrimaryOutputType() !=
+      B->getOutput().getPrimaryOutputType())
+    return false;
+  return A->getOutput().hasSameAdditionalOutputTypes(B->getOutput());
+}
+
+static bool
+jobsHaveSameEnvironment(const Job *A, const Job *B) {
+  auto AEnv = A->getExtraEnvironment();
+  auto BEnv = B->getExtraEnvironment();
+  if (AEnv.size() != BEnv.size())
+    return false;
+  for (size_t i = 0; i < AEnv.size(); ++i) {
+    if (strcmp(AEnv[i].first, BEnv[i].first) != 0)
+      return false;
+    if (strcmp(AEnv[i].second, BEnv[i].second) != 0)
+      return false;
+  }
+  return true;
+}
+
+bool
+ToolChain::jobIsBatchable(const Compilation &C, const Job *A) const {
+  // FIXME: There might be a tighter criterion to use here?
+  if (C.getOutputInfo().CompilerMode != OutputInfo::Mode::StandardCompile)
+    return false;
+  auto const *CJActA = dyn_cast<const CompileJobAction>(&A->getSource());
+  if (!CJActA)
+    return false;
+  return findSingleSwiftInput(CJActA) != nullptr;
+}
+
+bool
+ToolChain::jobsAreBatchCombinable(const Compilation &C,
+                                  const Job *A, const Job *B) const {
+  assert(jobIsBatchable(C, A));
+  assert(jobIsBatchable(C, B));
+  return (jobsHaveSameExecutableNames(A, B) &&
+          jobsHaveSameOutputTypes(A, B) &&
+          jobsHaveSameEnvironment(A, B));
+}
+
+/// Form a synthetic \c CommandOutput for a \c BatchJob by merging together the
+/// \c CommandOutputs of all the jobs passed.
+static std::unique_ptr<CommandOutput>
+makeBatchCommandOutput(ArrayRef<const Job *> jobs, Compilation &C,
+                       types::ID outputType) {
+  auto output =
+      llvm::make_unique<CommandOutput>(outputType, C.getDerivedOutputFileMap());
+  for (auto const *J : jobs) {
+    output->addOutputs(J->getOutput());
+  }
+  return output;
+}
+
+/// Set-union the \c Inputs and \c Actions from each \c Job in \p jobs into the
+/// provided \p inputJobs and \p inputActions vectors, further adding all \c
+/// Actions from the resulting merger to \p batchCJA. Do set-union rather than
+/// concatenation here to avoid mentioning the same input multiple times.
+static bool
+mergeBatchInputs(ArrayRef<const Job *> jobs,
+                 llvm::SmallSetVector<const Job *, 16> &inputJobs,
+                 llvm::SmallSetVector<const Action *, 16> &inputActions,
+                 CompileJobAction *batchCJA) {
+  for (auto const *J : jobs) {
+    for (auto const *I : J->getInputs()) {
+      inputJobs.insert(I);
+    }
+    auto const *CJA = dyn_cast<CompileJobAction>(&J->getSource());
+    if (!CJA)
+      return true;
+    for (auto const *I : CJA->getInputs()) {
+      inputActions.insert(I);
+    }
+  }
+
+  for (auto const *I : inputActions) {
+    batchCJA->addInput(I);
+  }
+  return false;
+}
+
+/// Construct a \c BatchJob by merging the constituent \p jobs' CommandOutput,
+/// input \c Job and \c Action members. Call through to \c constructInvocation
+/// on \p BatchJob, to build the \c InvocationInfo.
+std::unique_ptr<Job>
+ToolChain::constructBatchJob(ArrayRef<const Job *> jobs,
+                             Compilation &C) const
+{
+#ifndef NDEBUG
+  // Verify that the equivalence relation on the jobs also holds pairwise.
+  for (auto *A : jobs) {
+    for (auto *B : jobs) {
+      assert(jobsAreBatchCombinable(C, A, B));
+    }
+  }
+#endif
+  if (jobs.size() == 0)
+    return nullptr;
+
+  // Synthetic OutputInfo is a slightly-modified version of the initial
+  // compilation's OI.
+  auto OI = C.getOutputInfo();
+  OI.CompilerMode = OutputInfo::Mode::BatchModeCompile;
+
+  auto const *executablePath = jobs[0]->getExecutable();
+  auto outputType = jobs[0]->getOutput().getPrimaryOutputType();
+  auto output = makeBatchCommandOutput(jobs, C, outputType);
+
+  llvm::SmallSetVector<const Job *, 16> inputJobs;
+  llvm::SmallSetVector<const Action *, 16> inputActions;
+  auto *batchCJA = C.createAction<CompileJobAction>(outputType);
+  if (mergeBatchInputs(jobs, inputJobs, inputActions, batchCJA))
+    return nullptr;
+
+  JobContext context{C, inputJobs.getArrayRef(), inputActions.getArrayRef(),
+                     *output, OI};
+  auto invocationInfo = constructInvocation(*batchCJA, context);
+  return llvm::make_unique<BatchJob>(*batchCJA,
+                                     inputJobs.takeVector(),
+                                     std::move(output),
+                                     executablePath,
+                                     std::move(invocationInfo.Arguments),
+                                     std::move(invocationInfo.ExtraEnvironment),
+                                     std::move(invocationInfo.FilelistInfos),
+                                     jobs);
+}
+
 bool
 ToolChain::sanitizerRuntimeLibExists(const ArgList &args,
                                      StringRef sanitizerName) const {