Add argument to parse headers via a source's flags

Sources/idt/idt.cc

@@ -12,12 +12,17 @@
 #include "llvm/ADT/SmallPtrSet.h"
 
 #include <algorithm>
+#include <array>
+#include <cstdint>
 #include <cstdlib>
-#include <iostream>
+#include <iterator>
+#include <memory>
 #include <optional>
+#include <set>
 #include <string>
 #include <tuple>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 namespace idt {
@@ -63,6 +68,12 @@ ignored_symbols("ignore",
                 llvm::cl::CommaSeparated,
                 llvm::cl::cat(idt::category));
 
+llvm::cl::opt<std::string>
+    cdb_source("cdb-source",
+               llvm::cl::desc("Path of a translation unit with an entry in the "
+                              "compilation database."),
+               llvm::cl::value_desc("path"), llvm::cl::cat(idt::category));
+
 template <typename Key, typename Compare, typename Allocator>
 bool contains(const std::set<Key, Compare, Allocator>& set, const Key& key) {
   return set.find(key) != set.end();
@@ -271,6 +282,16 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     return false;
   }
 
+  // When `--cdb-source` is set, the positional argument is parsed as the
+  // translation unit's main file; restrict fixits to it so we don't
+  // accidentally rewrite decls in transitively-included headers. Outside
+  // of `--cdb-source` mode, this is a no-op.
+  template <typename Decl_> bool is_in_main_tu(const Decl_ *D) const {
+    if (cdb_source.empty())
+      return true;
+    return source_manager_.isInMainFile(get_location(D));
+  }
+
   template <typename Decl_>
   inline bool is_in_system_header(const Decl_ *D) const {
     return source_manager_.isInSystemHeader(get_location(D));
@@ -347,6 +368,10 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     if (!is_in_header(FD))
       return;
 
+    // Restrict to the main TU when `--cdb-source` is set.
+    if (!is_in_main_tu(FD))
+      return;
+
     // Ignore friend declarations.
     if (FD->getFriendObjectKind() != clang::Decl::FOK_None)
       return;
@@ -431,6 +456,10 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     if (!is_in_header(VD))
       return;
 
+    // Restrict to the main TU when `--cdb-source` is set.
+    if (!is_in_main_tu(VD))
+      return;
+
     // Skip local variables. We are only interested in static fields.
     if (VD->getParentFunctionOrMethod())
       return;
@@ -501,6 +530,10 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     if (is_in_system_header(RD))
       return;
 
+    // Restrict to the main TU when `--cdb-source` is set.
+    if (!is_in_main_tu(RD))
+      return;
+
     // Skip exporting template classes. For fully-specialized template classes,
     // isTemplated() returns false so they will be annotated if needed.
     if (RD->isTemplated())
@@ -737,6 +770,74 @@ struct factory : clang::tooling::FrontendActionFactory {
     return std::make_unique<idt::action>();
   }
 };
+
+// CompilationDatabase wrapper that looks up flags under one path (the
+// "donor") and reports them as belonging to a different path (the actual
+// input). Used by `--cdb-source` for projects where headers do not appear in
+// the compilation database.
+class cdb_source_database : public clang::tooling::CompilationDatabase {
+public:
+  cdb_source_database(const clang::tooling::CompilationDatabase &inner,
+                      std::string donor)
+      : inner_(inner), donor_(std::move(donor)) {}
+
+  std::vector<clang::tooling::CompileCommand>
+  getCompileCommands(llvm::StringRef path) const override {
+    auto commands = inner_.getCompileCommands(donor_);
+    for (auto &command : commands)
+      rewrite(command, path);
+    return commands;
+  }
+
+  // Stub these out for safety since we do not use them.
+  std::vector<std::string> getAllFiles() const override { return {}; }
+  std::vector<clang::tooling::CompileCommand>
+  getAllCompileCommands() const override {
+    return {};
+  }
+
+private:
+  static void rewrite(clang::tooling::CompileCommand &command,
+                      llvm::StringRef path) {
+    static constexpr std::array<llvm::StringRef, 5> kSourceExts{
+        ".cpp", ".cc", ".cxx", ".c++", ".C"};
+
+    // Replace the first source-file argument with `path`. If none is
+    // present, append `path` as the input.
+    bool swapped = false;
+    for (auto &arg : command.CommandLine) {
+      if (swapped)
+        break;
+      llvm::StringRef ref{arg};
+      for (llvm::StringRef ext : kSourceExts) {
+        if (ref.ends_with(ext)) {
+          arg = path.str();
+          swapped = true;
+          break;
+        }
+      }
+    }
+    if (!swapped)
+      command.CommandLine.push_back(path.str());
+
+    // Enforce `-x c++-header`. Replace any existing `-x <lang>` value, or
+    // insert `-x c++-header` immediately after argv[0].
+    auto x_it =
+        std::find(command.CommandLine.begin(), command.CommandLine.end(), "-x");
+    if (x_it != command.CommandLine.end() &&
+        std::next(x_it) != command.CommandLine.end()) {
+      *std::next(x_it) = "c++-header";
+    } else if (!command.CommandLine.empty()) {
+      command.CommandLine.insert(std::next(command.CommandLine.begin()),
+                                 {"-x", "c++-header"});
+    }
+
+    command.Filename = path.str();
+  }
+
+  const clang::tooling::CompilationDatabase &inner_;
+  std::string donor_;
+};
 }
 
 int main(int argc, char *argv[]) {
@@ -746,7 +847,13 @@ int main(int argc, char *argv[]) {
       CommonOptionsParser::create(argc, const_cast<const char **>(argv),
                                   idt::category, llvm::cl::OneOrMore);
   if (options) {
-    ClangTool tool{options->getCompilations(), options->getSourcePathList()};
+    std::unique_ptr<CompilationDatabase> wrapped;
+    CompilationDatabase *cdb = &options->getCompilations();
+    if (!cdb_source.empty()) {
+      wrapped = std::make_unique<idt::cdb_source_database>(*cdb, cdb_source);
+      cdb = wrapped.get();
+    }
+    ClangTool tool{*cdb, options->getSourcePathList()};
     return tool.run(new idt::factory{});
   } else {
     llvm::logAllUnhandledErrors(std::move(options.takeError()), llvm::errs());

Tests/CDBSource.cc

@@ -0,0 +1,5 @@
+// RUN: rm -rf %t && mkdir %t
+// RUN: echo "[{\"directory\":\"%p\",\"command\":\"clang++ -c %s\",\"file\":\"%s\"}]" | sed -e 's/\\/\\\\/g' > %t/compile_commands.json
+// RUN: %idt -p %t --cdb-source %s -export-macro IDT_TEST_ABI %S/include/CDBSource.h 2>&1 | %FileCheck %s
+
+// CHECK: CDBSource.h:1:1: remark: unexported public interface 'cdb_source_target'

Tests/CDBSourceMainTU.cc

@@ -0,0 +1,9 @@
+// RUN: rm -rf %t && mkdir %t
+// RUN: echo "[{\"directory\":\"%p\",\"command\":\"clang++ -c %s\",\"file\":\"%s\"}]" | sed -e 's/\\/\\\\/g' > %t/compile_commands.json
+// RUN: %idt -p %t --cdb-source %s -export-macro IDT_TEST_ABI %S/include/CDBSourceMainTU.h 2>&1 | %FileCheck %s
+
+// The outer header's decl is in the main TU and gets a fixit.
+// CHECK:     CDBSourceMainTU.h:2:1: remark: unexported public interface 'outer_decl'
+// The inner header is transitively included but is *not* the main file,
+// so its decl must be skipped.
+// CHECK-NOT: CDBSourceMainTUInner.h

Tests/include/CDBSource.h

@@ -0,0 +1 @@
+int cdb_source_target(int);

Tests/include/CDBSourceMainTU.h

@@ -0,0 +1,2 @@
+#include "CDBSourceMainTUInner.h"
+int outer_decl(int);

Tests/include/CDBSourceMainTUInner.h

@@ -0,0 +1 @@
+int inner_decl(int);