More closely mimic the Clang compilation

toolchain/base/BUILD

@@ -24,6 +24,21 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "in_flight_clang",
+    hdrs = ["in_flight_clang.h"],
+    srcs = ["in_flight_clang.cpp"],
+    deps = [
+        "//common:check",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//clang:ast",
+        "@llvm-project//clang:basic",
+        "@llvm-project//clang:driver",
+        "@llvm-project//clang:frontend",
+        "@llvm-project//clang:frontend_tool",
+    ],
+)
+
 cc_library(
     name = "kind_switch",
     hdrs = ["kind_switch.h"],

toolchain/base/in_flight_clang.cpp

@@ -0,0 +1,200 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "toolchain/base/in_flight_clang.h"
+
+#include "clang/Basic/Stack.h"
+#include "clang/Driver/Compilation.h"
+#include "clang/Driver/Driver.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/FrontendAction.h"
+#include "clang/Frontend/MultiplexConsumer.h"
+#include "clang/Frontend/TextDiagnostic.h"
+#include "clang/Sema/Lookup.h"
+#include "clang/include/clang/FrontendTool/Utils.h"
+#include "common/check.h"
+
+namespace Carbon {
+
+struct InFlightClang::AstData {
+  std::mutex mut;
+  std::condition_variable cond;
+
+  clang::Sema* sema = nullptr;
+  // TODO: it is essentially a state machine, switch to an enum.
+  bool ast_ready = false;
+  bool finish_compilation_requested = false;
+  bool worker_finished = false;
+};
+
+InFlightClang::InFlightClang(clang::Sema* sema, std::unique_ptr<AstData> data,
+                             std::thread worker)
+    : sema_(sema), data_(std::move(data)), worker_(std::move(worker)) {}
+
+auto InFlightClang::getASTContext() -> clang::ASTContext& {
+  return getSema().getASTContext();
+}
+
+auto InFlightClang::getSourceManager() -> clang::SourceManager& {
+  return getSema().getSourceManager();
+}
+
+auto InFlightClang::getSourceManager() const -> const clang::SourceManager& {
+  CARBON_CHECK(sema_ != nullptr);
+  return sema_->getSourceManager();
+}
+
+auto InFlightClang::getSema() -> clang::Sema& {
+  CARBON_CHECK(sema_ != nullptr);
+  return *sema_;
+}
+
+auto InFlightClang::CompileFromArguments(
+    llvm::ArrayRef<const char*> argv, llvm::StringRef target,
+    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fs,
+    std::unique_ptr<clang::DiagnosticConsumer> consumer)
+    -> std::unique_ptr<InFlightClang> {
+  // Prepare the compiler run.
+  auto compiler = std::make_unique<clang::CompilerInstance>();
+  compiler->createDiagnostics(*fs, consumer.release(), /*ShouldOwn=*/true);
+
+  clang::driver::Driver driver(argv[0], target, compiler->getDiagnostics(),
+                               "clang LLVM compiler", fs);
+  std::unique_ptr<clang::driver::Compilation> compilation(
+      driver.BuildCompilation(argv));
+  if (!compilation) {
+    return nullptr;
+  }
+  const clang::driver::JobList& jobs = compilation->getJobs();
+  if (jobs.size() != 1) {
+    // CARBON_VLOG("got multiple jobs with '{}'", args);
+    return nullptr;
+  }
+  auto& job = *jobs.begin();
+  if (job.getSource().getKind() != clang::driver::Action::AssembleJobClass) {
+    // CARBON_VLOG("expected an assemble job class with '{}'", args);k
+    return nullptr;
+  }
+  if (!clang::CompilerInvocation::CreateFromArgs(
+          compiler->getInvocation(), job.getArguments(),
+          compiler->getDiagnostics(), argv[0])) {
+    return nullptr;
+  }
+  class InFlightClangASTConsumer : public clang::SemaConsumer {
+   public:
+    explicit InFlightClangASTConsumer(AstData* data)
+        : data_(data), active_sema_(nullptr) {}
+
+    void InitializeSema(clang::Sema& sema) override { active_sema_ = &sema; }
+    void ForgetSema() override { active_sema_ = nullptr; }
+
+    void HandleTranslationUnit(clang::ASTContext& ast) override {
+      (void)ast;
+      // Provide the AST to the consumers.
+      {
+        std::unique_lock<std::mutex> lock(data_->mut);
+        CARBON_CHECK(data_->sema == nullptr);
+        data_->sema = active_sema_;
+        data_->ast_ready = true;
+
+        lock.unlock();
+        data_->cond.notify_all();
+      }
+      // And wait until we are asked to finish the compilation.
+      {
+        std::unique_lock<std::mutex> lock(data_->mut);
+        data_->cond.wait(lock,
+                         [&] { return data_->finish_compilation_requested; });
+      }
+    }
+
+   private:
+    AstData* data_;
+    clang::Sema* active_sema_;
+  };
+
+  class InFlightClangFrontendAction : public clang::WrapperFrontendAction {
+   public:
+    explicit InFlightClangFrontendAction(
+        AstData* data, std::unique_ptr<FrontendAction> wrapped)
+        : clang::WrapperFrontendAction(std::move(wrapped)), data_(data) {}
+
+    auto CreateASTConsumer(clang::CompilerInstance& compiler,
+                           llvm::StringRef file)
+        -> std::unique_ptr<clang::ASTConsumer> override {
+      std::vector<std::unique_ptr<clang::ASTConsumer>> consumers;
+      consumers.push_back(std::make_unique<InFlightClangASTConsumer>(data_));
+      if (auto c = WrapperFrontendAction::CreateASTConsumer(compiler, file)) {
+        consumers.push_back(std::move(c));
+      }
+
+      return std::make_unique<clang::MultiplexConsumer>(std::move(consumers));
+    }
+
+   private:
+    AstData* data_;
+  };
+
+  auto data = std::make_unique<AstData>();
+  auto action = std::make_unique<InFlightClangFrontendAction>(
+      data.get(), clang::CreateFrontendAction(*compiler));
+  if (!action->PrepareToExecute(*compiler)) {
+    return nullptr;
+  }
+  if (!action->BeginSourceFile(*compiler,
+                               compiler->getFrontendOpts().Inputs.front())) {
+    return nullptr;
+  }
+
+  // Clang does not have the API to get the compilation to a point where the C++
+  // translation unit is finished, but lowering (CodeGen in Clang) has not
+  // started yet.
+  // It does provide us callbacks and will run on a separate thread to emulate
+  // this API.
+  std::thread worker([compiler = std::move(compiler),
+                      action = std::move(action), data = data.get()] {
+    // TODO: ensure the created thread gets a large stack.
+    clang::noteBottomOfStack();
+
+    // Note: the ASTConsumer will be called internally and communicate the
+    // AST back to the clients.
+    auto err = action->Execute();
+    llvm::cantFail(std::move(err), "TODO: log this error");
+
+    {
+      std::unique_lock<std::mutex> lock(data->mut);
+      CARBON_CHECK(!data->worker_finished);
+      data->worker_finished = true;
+    }
+    data->cond.notify_all();
+  });
+
+  // Wait until the AST is ready, we use threads to transform a callback-based
+  // API into a simpler style.
+  clang::Sema* sema = nullptr;
+  {
+    std::unique_lock<std::mutex> lock(data->mut);
+    data->cond.wait(lock,
+                    [&] { return data->ast_ready || data->worker_finished; });
+    sema = data->sema;
+  }
+  // The worker thread is waiting for a signal to finish the compilation at this
+  // moment. We will signal it to continue in InFlightClang's destructor, at a
+  // point where nobody should be using the AST anymore.
+
+  return std::unique_ptr<InFlightClang>(
+      new InFlightClang(sema, std::move(data), std::move(worker)));
+}
+
+InFlightClang::~InFlightClang() {
+  {
+    std::unique_lock<std::mutex> lock(data_->mut);
+    data_->finish_compilation_requested = true;
+  }
+  data_->cond.notify_all();
+
+  worker_.join();
+}
+
+}  // namespace Carbon

toolchain/base/in_flight_clang.h

@@ -0,0 +1,48 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef CARBON_TOOLCHAIN_BASE_IN_FLIGHT_CLANG_H
+#define CARBON_TOOLCHAIN_BASE_IN_FLIGHT_CLANG_H
+
+#include <thread>
+
+#include "clang/AST/ASTContext.h"
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/SourceManager.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
+
+namespace Carbon {
+class InFlightClang {
+ public:
+  ~InFlightClang();
+
+  // Runs the compiler on the passed code and stops it at a point suited for
+  // doing any additional operations on the frontend.
+  //
+  // The arguments must produce exactly one compile job.
+  static auto CompileFromArguments(
+      llvm::ArrayRef<const char*> argv, llvm::StringRef target,
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fs,
+      std::unique_ptr<clang::DiagnosticConsumer> consumer)
+      -> std::unique_ptr<InFlightClang>;
+
+  auto getASTContext() -> clang::ASTContext&;
+  auto getSourceManager() -> clang::SourceManager&;
+  auto getSourceManager() const -> const clang::SourceManager&;
+  auto getSema() -> clang::Sema&;
+
+ private:
+  struct AstData;
+  InFlightClang(clang::Sema* sema, std::unique_ptr<AstData> data,
+                std::thread worker);
+
+  clang::Sema* const sema_ = nullptr;
+  std::unique_ptr<AstData> data_;
+  std::thread worker_;
+};
+
+}  // namespace Carbon
+
+#endif

toolchain/check/check.h

@@ -6,6 +6,7 @@
 #define CARBON_TOOLCHAIN_CHECK_CHECK_H_
 
 #include "common/ostream.h"
+#include "toolchain/base/in_flight_clang.h"
 #include "toolchain/base/shared_value_stores.h"
 #include "toolchain/base/timings.h"
 #include "toolchain/check/diagnostic_emitter.h"
@@ -26,7 +27,7 @@ struct Unit {
   SemIR::File* sem_ir;
 
   // The Clang AST owned by `CompileSubcommand`.
-  std::unique_ptr<clang::ASTUnit>* cpp_ast = nullptr;
+  std::unique_ptr<InFlightClang>* cpp_ast = nullptr;
 };
 
 // Checks a group of parse trees. This will use imports to decide the order of