[include-fixer] Add the missing header to the file where the unidentified symbol comes from.

[bkramer]

include-fixer/IncludeFixer.cpp

104

Does this patch do the right thing for the test case in the comment? Otherwise we'll try add includes to system headers all the time, which both wrong and invisible to the user.

[hokein]

include-fixer/IncludeFixer.cpp

104

No, the patch doesn't handle this case :(. Currently it is hard to figure out this case. For the system headers, we can probably blacklist them?

include-fixer/IncludeFixer.cpp

 public:
   explicit Action(SymbolIndexManager &SymbolIndexMgr, bool MinimizeIncludePaths)
       : SymbolIndexMgr(SymbolIndexMgr),
         MinimizeIncludePaths(MinimizeIncludePaths) {}
 
   std::unique_ptr<clang::ASTConsumer>
   CreateASTConsumer(clang::CompilerInstance &Compiler,
                     StringRef InFile) override {
-    Filename = InFile;
     return llvm::make_unique<clang::ASTConsumer>();
   }
 
   void ExecuteAction() override {
     clang::CompilerInstance *Compiler = &getCompilerInstance();
     assert(!Compiler->hasSema() && "CI already has Sema");
 
     // Set up our hooks into sema and parse the AST.
clang::TypoCorrection CorrectTypo(const DeclarationNameInfo &Typo,
                                     CorrectionCandidateCallback &CCC,
                                     DeclContext *MemberContext,
                                     bool EnteringContext,
                                     const ObjCObjectPointerType *OPT) override {
     // Ignore spurious callbacks from SFINAE contexts.
     if (getCompilerInstance().getSema().isSFINAEContext())
       return clang::TypoCorrection();
 
-    // We currently ignore the unidentified symbol which is not from the
-    // main file.
-    //
-    // However, this is not always true due to templates in a non-self contained
-    // header, consider the case:
-    //
-    //   // header.h
-    //   template <typename T>
-    //   class Foo {
-    //     T t;
-    //   };
-    //
-    //   // test.cc
-    //   // We need to add <bar.h> in test.cc instead of header.h.
-    //   class Bar;
-    //   Foo<Bar> foo;
-    //
-    // FIXME: Add the missing header to the header file where the symbol comes
-    // from.
-    if (!getCompilerInstance().getSourceManager().isWrittenInMainFile(
-            Typo.getLoc()))
-      return clang::TypoCorrection();
-
     std::string TypoScopeString;
     if (S) {
       // FIXME: Currently we only use namespace contexts. Use other context
       // types for query.
       for (const auto *Context = S->getEntity(); Context;
            Context = Context->getParent()) {
         if (const auto *ND = dyn_cast<NamespaceDecl>(Context)) {
           if (!ND->getName().empty())
if (!query(TypoScopeString + QueryString, Typo.getLoc()))
       query(QueryString, Typo.getLoc());
 
     // FIXME: We should just return the name we got as input here and prevent
     // clang from trying to correct the typo by itself. That may change the
     // identifier to something that's not wanted by the user.
     return clang::TypoCorrection();
   }
 
-  StringRef filename() const { return Filename; }
-
   /// Get the minimal include for a given path.
   std::string minimizeInclude(StringRef Include,
                               const clang::SourceManager &SourceManager,
                               clang::HeaderSearch &HeaderSearch) {
     if (!MinimizeIncludePaths)
       return Include;
 
     // Get the FileEntry for the include.
public:
   }
 
   /// Get the include fixer context for the queried symbol.
   IncludeFixerContext
   getIncludeFixerContext(const clang::SourceManager &SourceManager,
                          clang::HeaderSearch &HeaderSearch) {
     IncludeFixerContext FixerContext;
     FixerContext.SymbolIdentifier = QuerySymbol;
+    SmallString<256> CleanedFilePath = StringRef(Filename);
+    llvm::sys::path::remove_dots(CleanedFilePath, /*remove_dot_dot=*/true);
+    FixerContext.FilePath = CleanedFilePath.str();
     for (const auto &Header : SymbolQueryResults)
       FixerContext.Headers.push_back(
           minimizeInclude(Header, SourceManager, HeaderSearch));
 
     return FixerContext;
   }
 
 private:
   /// Query the database for a given identifier.
   bool query(StringRef Query, SourceLocation Loc) {
     assert(!Query.empty() && "Empty query!");
 
     // Skip other identifers once we have discovered an identfier successfully.
     if (!SymbolQueryResults.empty())
       return false;
 
     DEBUG(llvm::dbgs() << "Looking up '" << Query << "' at ");
     DEBUG(Loc.print(llvm::dbgs(), getCompilerInstance().getSourceManager()));
     DEBUG(llvm::dbgs() << " ...");
 
     QuerySymbol = Query.str();
+    Filename = getCompilerInstance().getSourceManager().getFilename(Loc);
     SymbolQueryResults = SymbolIndexMgr.search(Query);
     DEBUG(llvm::dbgs() << SymbolQueryResults.size() << " replies\n");
     return !SymbolQueryResults.empty();
   }
 
   /// The client to use to find cross-references.
   SymbolIndexManager &SymbolIndexMgr;
 
-  /// The absolute path to the file being processed.
+  /// The absolute path to a file where the QuerySymbol comes from.
   std::string Filename;
 
   /// The symbol being queried.
   std::string QuerySymbol;
 
   /// The query results of an identifier. We only include the first discovered
   /// identifier to avoid getting caught in results from error recovery.
   std::vector<std::string> SymbolQueryResults;

include-fixer/IncludeFixerContext.h

 
 namespace clang {
 namespace include_fixer {
 
 /// \brief A context for the symbol being queried.
 struct IncludeFixerContext {
   /// \brief The symbol name.
   std::string SymbolIdentifier;
+  /// \bridf The absolute path to a file where the SymbolIdentifier comes from.
+  std::string FilePath;
   /// \brief The headers which have SymbolIdentifier definitions.
   std::vector<std::string> Headers;
 };
 
 } // namespace include_fixer
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXERCONTEXT_H

include-fixer/tool/ClangIncludeFixer.cpp

 LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(std::string)
 
 namespace llvm {
 namespace yaml {
 template <> struct MappingTraits<IncludeFixerContext> {
   static void mapping(IO &io, IncludeFixerContext &Context) {
     io.mapRequired("SymbolIdentifier", Context.SymbolIdentifier);
     io.mapRequired("Headers", Context.Headers);
+    io.mapRequired("FilePath", Context.FilePath);
   }
 };
 } // namespace yaml
 } // namespace llvm
 
 namespace {
 cl::OptionCategory IncludeFixerCategory("Tool options");
 
STDINMode("stdin",
               cl::init(false), cl::cat(IncludeFixerCategory));
 
 cl::opt<bool> OutputHeaders(
     "output-headers",
     cl::desc("Print the symbol being queried and all its relevant headers in\n"
              "JSON format to stdout:\n"
              "  {\n"
              "    \"SymbolIdentifier\": \"foo\",\n"
+             "    \"FilePath\": \"path/to/main.cc\", \n"
              "    \"Headers\": [\"\\\"foo_a.h\\\"\"]\n"
              "  }"),
     cl::init(false), cl::cat(IncludeFixerCategory));
 
 cl::opt<std::string> InsertHeader(
     "insert-header",
     cl::desc("Insert a specific header. This should run with STDIN mode.\n"
              "The result is written to stdout. It is currently used for\n"
              "editor integration. Support YAML/JSON format:\n"
              "  -insert-header=\"{SymbolIdentifier: foo,\n"
+             "                   FilePath: path/to/main.cc,\n"
              "                   Headers: ['\\\"foo_a.h\\\"']}\""),
     cl::init(""), cl::cat(IncludeFixerCategory));
 
 cl::opt<std::string>
     Style("style",
           cl::desc("Fallback style for reformatting after inserting new "
                    "headers if there is no clang-format config file found."),
           cl::init("llvm"), cl::cat(IncludeFixerCategory));
createSymbolIndexManager(StringRef FilePath) {
   }
   }
   return SymbolIndexMgr;
 }
 
 void writeToJson(llvm::raw_ostream &OS, const IncludeFixerContext& Context) {
   OS << "{\n"
         "  \"SymbolIdentifier\": \"" << Context.SymbolIdentifier << "\",\n"
+        "  \"FilePath\": \"" << Context.FilePath << "\",\n"
         "  \"Headers\": [ ";
   for (const auto &Header : Context.Headers) {
     OS << " \"" << llvm::yaml::escape(Header) << "\"";
     if (Header != Context.Headers.back())
       OS << ", ";
   }
   OS << " ]\n"
         "}\n";
if (STDINMode) {
     }
     Code = std::move(CodeOrErr.get());
     if (Code->getBufferSize() == 0)
       return 0;  // Skip empty files.
 
     tool.mapVirtualFile(options.getSourcePathList().front(), Code->getBuffer());
   }
 
-  StringRef FilePath = options.getSourcePathList().front();
-  format::FormatStyle InsertStyle = format::getStyle("file", FilePath, Style);
+  format::FormatStyle InsertStyle =
+      format::getStyle("file", options.getSourcePathList().front(), Style);
 
   if (!InsertHeader.empty()) {
     if (!STDINMode) {
       errs() << "Should be running in STDIN mode\n";
       return 1;
     }
 
     llvm::yaml::Input yin(InsertHeader);
     IncludeFixerContext Context;
     yin >> Context;
 
     if (Context.Headers.size() != 1) {
       errs() << "Expect exactly one inserted header.\n";
       return 1;
     }
 
     tooling::Replacements Replacements =
         clang::include_fixer::createInsertHeaderReplacements(
-            Code->getBuffer(), FilePath, Context.Headers[0], InsertStyle);
+            Code->getBuffer(), Context.FilePath, Context.Headers[0],
+            InsertStyle);
     tooling::Replacements Replaces(Replacements.begin(), Replacements.end());
     std::string ChangedCode =
         tooling::applyAllReplacements(Code->getBuffer(), Replaces);
     llvm::outs() << ChangedCode;
     return 0;
   }
 
   // Set up data source.
int includeFixerMain(int argc, const char **argv) {
   if (OutputHeaders) {
     writeToJson(llvm::outs(), Context);
     return 0;
   }
 
   if (Context.Headers.empty())
     return 0;
 
-  auto Buffer = llvm::MemoryBuffer::getFile(FilePath);
+  auto Buffer = llvm::MemoryBuffer::getFile(Context.FilePath);
   if (!Buffer) {
-    errs() << "Couldn't open file: " << FilePath;
+    errs() << "Couldn't open file: " << Context.FilePath;
     return 1;
   }
 
-  // FIXME: Rank the results and pick the best one instead of the first one.
   tooling::Replacements Replacements =
       clang::include_fixer::createInsertHeaderReplacements(
-          /*Code=*/Buffer.get()->getBuffer(), FilePath, Context.Headers.front(),
-          InsertStyle);
+          /*Code=*/Buffer.get()->getBuffer(), Context.FilePath,
+          Context.Headers.front(), InsertStyle);
 
   if (!Quiet)
     llvm::errs() << "Added #include" << Context.Headers.front();
 
   // Set up a new source manager for applying the resulting replacements.
   IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts(new DiagnosticOptions);
   DiagnosticsEngine Diagnostics(new DiagnosticIDs, &*DiagOpts);
   TextDiagnosticPrinter DiagnosticPrinter(outs(), &*DiagOpts);

include-fixer/tool/clang-include-fixer.py

 
 def execute(command, text):
   p = subprocess.Popen(command,
                        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                        stdin=subprocess.PIPE)
   return p.communicate(input=text)
 
 
-def InsertHeaderToVimBuffer(header, text):
+def InsertHeaderToVimBuffer(header, code):
   command = [binary, "-stdin", "-insert-header="+json.dumps(header),
              vim.current.buffer.name]
-  stdout, stderr = execute(command, text)
+  stdout, stderr = execute(command, code)
+
   if stdout:
+    # If the file being inserted is not the current one in vim buffer, just
+    # overwrite it in filesystem.
+    if header["FilePath"] != vim.current.buffer.name:
+      with open(header["FilePath"], "w") as f:
+        f.write(stdout)
+      return
+
     lines = stdout.splitlines()
     sequence = difflib.SequenceMatcher(None, vim.current.buffer, lines)
     for op in reversed(sequence.get_opcodes()):
       if op[0] is not 'equal':
         vim.current.buffer[op[1]:op[2]] = lines[op[3]:op[4]]
 
 
 def main():
   parser = argparse.ArgumentParser(
       description='Vim integration for clang-include-fixer')
   parser.add_argument('-db', default='yaml',
                       help='clang-include-fixer input format.')
   parser.add_argument('-input', default='',
                       help='String to initialize the database.')
   args = parser.parse_args()
 
   # Get the current text.
   buf = vim.current.buffer
-  text = '\n'.join(buf)
+  code = '\n'.join(buf)
 
   # Run command to get all headers.
   command = [binary, "-stdin", "-output-headers", "-db="+args.db,
              "-input="+args.input, vim.current.buffer.name]
-  stdout, stderr = execute(command, text)
+  stdout, stderr = execute(command, code)
   if stderr:
     print >> sys.stderr, "Error while running clang-include-fixer: " + stderr
     return
 
   include_fixer_context = json.loads(stdout)
   symbol = include_fixer_context["SymbolIdentifier"]
   headers = include_fixer_context["Headers"]
+  file_path = include_fixer_context["FilePath"]
+
+  if file_path != vim.current.buffer.name:
+    with open(file_path, "r") as f:
+      code = f.read();
 
   if not symbol:
     print "The file is fine, no need to add a header.\n"
     return;
 
   if not headers:
     print "Couldn't find a header for {0}.\n".format(symbol)
     return
 
   # The first line is the symbol name.
   # If there is only one suggested header, insert it directly.
   if len(headers) == 1 or maximum_suggested_headers == 1:
     InsertHeaderToVimBuffer({"SymbolIdentifier": symbol,
-                             "Headers":[headers[0]]}, text)
-    print "Added #include {0} for {1}.\n".format(headers[0], symbol)
+                             "Headers": [headers[0]],
+                             "FilePath": file_path}, code)
+    print "Added #include {0} for {1} in {2}.\n".format(headers[0],
+                                                        symbol,
+                                                        file_path)
     return
 
   choices_message = ""
   index = 1;
   for header in headers[0:maximum_suggested_headers]:
     choices_message += "&{0} {1}\n".format(index, header)
     index += 1
 
   select = ShowDialog("choose a header file for {0}.".format(symbol),
                       choices_message)
   # Insert a selected header.
   InsertHeaderToVimBuffer({"SymbolIdentifier": symbol,
-                           "Headers":[headers[select-1]]}, text)
-  print "Added #include {0} for {1}.\n".format(headers[select-1], symbol)
+                           "Headers": [headers[select-1]],
+                           "FilePath": file_path}, code)
+  print "Added #include {0} for {1} in {2}.\n".format(headers[select-1],
+                                                      symbol,
+                                                      file_path)
   return;
 
 
 if __name__ == '__main__':
   main()

test/include-fixer/commandline_options.cpp

 // REQUIRES: shell
 // RUN: sed -e 's#//.*$##' %s > %t.cpp
 // RUN: clang-include-fixer -db=fixed -input='foo= "foo.h","bar.h"' -output-headers %t.cpp -- | FileCheck %s -check-prefix=CHECK-HEADERS
-// RUN: cat %t.cpp | clang-include-fixer -stdin -insert-header='{SymbolIdentifier: foo, Headers: ["\"foo.h\""]}' %t.cpp | FileCheck %s -check-prefix=CHECK
+// RUN: cat %t.cpp | clang-include-fixer -stdin -insert-header='{SymbolIdentifier: foo, FilePath: %t.cpp, Headers: ["\"foo.h\""]}' %t.cpp | FileCheck %s -check-prefix=CHECK
 //
 // CHECK-HEADERS: "Headers": [ "\"foo.h\"", "\"bar.h\"" ]
 //
 // CHECK: #include "foo.h"
 // CHECK: foo f;
 
 foo f;

test/include-fixer/insert_header_in_header.cpp

+// REQUIRES: shell
+// RUN: echo 'foo f;' > %T/main.h
+// RUN: sed -e 's#//.*$##' %s > %t.cpp
+// RUN: clang-include-fixer -db=yaml -input=%p/Inputs/fake_yaml_db.yaml %t.cpp --
+// RUN: FileCheck %s -input-file=%T/main.h
+
+// CHECK: #include "foo.h"
+// CHECK: foo f;
+
+#include "main.h"

unittests/include-fixer/IncludeFixerTest.cpp

InMemoryFileSystem->addFile(FileName, 0,
                               llvm::MemoryBuffer::getMemBuffer(Code));
 
   InMemoryFileSystem->addFile("foo.h", 0,
                               llvm::MemoryBuffer::getMemBuffer("\n"));
   InMemoryFileSystem->addFile("dir/bar.h", 0,
                               llvm::MemoryBuffer::getMemBuffer("\n"));
   InMemoryFileSystem->addFile("dir/otherdir/qux.h", 0,
                               llvm::MemoryBuffer::getMemBuffer("\n"));
-  InMemoryFileSystem->addFile("header.h", 0,
-                              llvm::MemoryBuffer::getMemBuffer("bar b;"));
   return Invocation.run();
 }
 
 static std::string runIncludeFixer(
     StringRef Code,
     const std::vector<std::string> &ExtraArgs = std::vector<std::string>()) {
   std::vector<SymbolInfo> Symbols = {
       SymbolInfo("string", SymbolInfo::SymbolKind::Class, "<string>", 1,
EXPECT_EQ("#include \"bar.h\"\nnamespace A {\nb::bar b;\n}",
             runIncludeFixer("namespace A {\nb::bar b;\n}"));
 }
 
 TEST(IncludeFixer, EnumConstantSymbols) {
   EXPECT_EQ("#include \"color.h\"\nint test = a::b::Green;\n",
             runIncludeFixer("int test = a::b::Green;\n"));
 }
 
-TEST(IncludeFixer, IgnoreSymbolFromHeader) {
-  std::string Code = "#include \"header.h\"";
-  EXPECT_EQ(Code, runIncludeFixer(Code));
-}
-
 // FIXME: add test cases for inserting and sorting multiple headers when
 // include-fixer supports multiple headers insertion.
 TEST(IncludeFixer, InsertAndSortSingleHeader) {
   // Insert one header.
   std::string Code = "#include \"a.h\"\n"
                      "#include \"foo.h\"\n"
                      "\n"
                      "namespace a { b::bar b; }";