Add argument to parse headers via a source's flags

Sources/idt/idt.cc

@@ -12,12 +12,16 @@
 #include "llvm/ADT/SmallPtrSet.h"
 
 #include <algorithm>
+#include <array>
+#include <cstdint>
 #include <cstdlib>
-#include <iostream>
+#include <memory>
 #include <optional>
+#include <set>
 #include <string>
 #include <tuple>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 namespace idt {
@@ -63,11 +67,23 @@ ignored_symbols("ignore",
                 llvm::cl::CommaSeparated,
                 llvm::cl::cat(idt::category));
 
+llvm::cl::opt<std::string>
+main_file("main-file",
+          llvm::cl::desc("Path of a translation unit with an entry in the "
+                         "compilation database whose flags will be used to "
+                         "parse the positional argument"),
+          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();
 }
 
+template <class T, std::size_t N>
+bool contains(const std::array<T, N>& arr, const T& value) {
+  return std::find(arr.begin(), arr.end(), value) != arr.end();
+}
+
 const std::set<std::string> &get_ignored_symbols() {
   static auto kIgnoredFunctions = [&]() -> std::set<std::string> {
       return { ignored_symbols.begin(), ignored_symbols.end() };
@@ -271,6 +287,16 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     return false;
   }
 
+  // When `--main-file` 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 `--main-file` mode, this is a no-op.
+  template <typename Decl_> bool is_in_main_tu(const Decl_ *D) const {
+    if (main_file.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 +373,10 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     if (!is_in_header(FD))
       return;
 
+    // Restrict to the main TU when `--main-file` is set.
+    if (!is_in_main_tu(FD))
+      return;
+
     // Ignore friend declarations.
     if (FD->getFriendObjectKind() != clang::Decl::FOK_None)
       return;
@@ -431,6 +461,10 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     if (!is_in_header(VD))
       return;
 
+    // Restrict to the main TU when `--main-file` 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 +535,10 @@ class visitor : public clang::RecursiveASTVisitor<visitor> {
     if (is_in_system_header(RD))
       return;
 
+    // Restrict to the main TU when `--main-file` 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 +775,59 @@ 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 `--main-file` for projects where headers do not appear in
+// the compilation database.
+class main_file_database : public clang::tooling::CompilationDatabase {
+public:
+  main_file_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 (contains(kSourceExts, llvm::sys::path::extension(llvm::StringRef(arg)))) {
+        arg = path.str();
+        swapped = true;
+        break;
+      }
+    }
+    if (!swapped)
+      command.CommandLine.push_back(path.str());
+
+    command.Filename = path.str();
+  }
+
+  const clang::tooling::CompilationDatabase &inner_;
+  std::string donor_;
+};
 }
 
 int main(int argc, char *argv[]) {
@@ -746,7 +837,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 (!main_file.empty()) {
+      wrapped = std::make_unique<idt::main_file_database>(*cdb, main_file);
+      cdb = wrapped.get();
+    }
+    ClangTool tool{*cdb, options->getSourcePathList()};
     return tool.run(new idt::factory{});
   } else {
     llvm::logAllUnhandledErrors(std::move(options.takeError()), llvm::errs());

Tests/MainFile.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 --main-file %s -export-macro IDT_TEST_ABI %S/include/MainFile.h 2>&1 | %FileCheck %s
+
+// CHECK: MainFile.h:1:1: remark: unexported public interface 'main_file_target'

Tests/MainFileMainTU.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 --main-file %s -export-macro IDT_TEST_ABI %S/include/MainFileMainTU.h 2>&1 | %FileCheck %s
+
+// The outer header's decl is in the main TU and gets a fixit.
+// CHECK:     MainFileMainTU.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: MainFileMainTUInner.h

Tests/include/MainFile.h

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

Tests/include/MainFileMainTU.h

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

Tests/include/MainFileMainTUInner.h

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