[clang][include cleaner] Fix MainHeader insertion issue (#212852) When doing multiple header insertions (like from include cleaner) there was an issue with MainHeaders inserted in the wrong location. HeaderIncludes now supports a bulk insertion where it sorts the insertions appropriately before inserting them to make sure that the MainHeader ends up in the correct location. Format.cpp has been updated to use the bulk insertion method correctly. Tests added to verify. GitOrigin-RevId: fc1ab1f288992fd22b4f562557ea38b0cc8f9aaf
diff --git a/include/clang/Tooling/Inclusions/HeaderIncludes.h b/include/clang/Tooling/Inclusions/HeaderIncludes.h index 72407e2..84b1d19 100644 --- a/include/clang/Tooling/Inclusions/HeaderIncludes.h +++ b/include/clang/Tooling/Inclusions/HeaderIncludes.h
@@ -9,13 +9,16 @@ #ifndef LLVM_CLANG_TOOLING_INCLUSIONS_HEADERINCLUDES_H #define LLVM_CLANG_TOOLING_INCLUSIONS_HEADERINCLUDES_H -#include "clang/Basic/SourceManager.h" #include "clang/Tooling/Core/Replacement.h" #include "clang/Tooling/Inclusions/IncludeStyle.h" -#include "llvm/Support/Path.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringMap.h" #include "llvm/Support/Regex.h" #include <list> #include <optional> +#include <set> +#include <string> #include <unordered_map> namespace clang { @@ -64,7 +67,7 @@ /// default. These code sections include: /// - raw string literals (containing #include). /// - #if blocks. - /// - Special #include's among declarations (e.g. functions). + /// - Special #includes among declarations (e.g. functions). /// /// Returns a replacement that inserts the new header into a suitable #include /// block of the same category. This respects the order of the existing @@ -77,6 +80,38 @@ bool IsAngled, IncludeDirective Directive) const; + /// Represents a single header directive to be inserted in a batch operation. + /// + /// Usage: + /// - HeaderToInsert("<vector>") -> inserts #include <vector> + /// (auto-detects angled) + /// - HeaderToInsert("\"foo.h\"") -> inserts #include "foo.h" + /// (auto-detects quoted) + /// - HeaderToInsert("<foo>", IncludeDirective::Import) -> inserts #import + /// <foo> + /// - HeaderToInsert("foo.h", IncludeDirective::Include, /*IsAngled=*/false) + /// -> explicit IsAngled + struct HeaderToInsert { + enum class QuoteStyle { AUTO, ANGLED, QUOTED }; + + // The header name, with any surrounding quotes or brackets removed. + std::string Header; + // Whether to insert #include or #import. + IncludeDirective Directive; + // Whether to use <> or "" for the header. This can be set explicitly with + // QuoteStyle::ANGLED or QuoteStyle::QUOTED, or auto-detected based on + // `RawOrSpelledHeader` with QuoteStyle::AUTO. + bool IsAngled; + + HeaderToInsert(llvm::StringRef RawOrSpelledHeader, + IncludeDirective Directive = IncludeDirective::Include, + QuoteStyle QuoteStyle = QuoteStyle::AUTO); + }; + + /// Inserts a batch of headers into the code, sorting and grouping them + /// according to IncludeStyle and returning the replacements. + tooling::Replacements insert(llvm::ArrayRef<HeaderToInsert> Headers) const; + /// Removes all existing #includes and #imports of \p Header quoted with <> if /// \p IsAngled is true or "" if \p IsAngled is false. /// This doesn't resolve the header file path; it only deletes #includes and @@ -88,7 +123,7 @@ private: struct Include { - Include(StringRef Name, tooling::Range R, IncludeDirective D) + Include(llvm::StringRef Name, tooling::Range R, IncludeDirective D) : Name(Name), R(R), Directive(D) {} // An include header quoted with either <> or "".
diff --git a/lib/Format/Format.cpp b/lib/Format/Format.cpp index 1923754..eec0a9d 100644 --- a/lib/Format/Format.cpp +++ b/lib/Format/Format.cpp
@@ -4217,6 +4217,7 @@ } SmallVector<StringRef, 4> Matches; + SmallVector<tooling::HeaderIncludes::HeaderToInsert, 4> HeadersToInsert; for (const auto &R : HeaderInsertions) { auto IncludeDirective = R.getReplacementText(); bool Matched = @@ -4225,19 +4226,17 @@ "'#include ...'"); (void)Matched; auto IncludeName = Matches[2]; - auto Replace = - Includes.insert(IncludeName.trim("\"<>"), IncludeName.starts_with("<"), - tooling::IncludeDirective::Include); - if (Replace) { - auto Err = Result.add(*Replace); - if (Err) { - consumeError(std::move(Err)); - unsigned NewOffset = - Result.getShiftedCodePosition(Replace->getOffset()); - auto Shifted = tooling::Replacement(FileName, NewOffset, 0, - Replace->getReplacementText()); - Result = Result.merge(tooling::Replacements(Shifted)); - } + HeadersToInsert.emplace_back(IncludeName, + tooling::IncludeDirective::Include); + } + for (const auto &Replace : Includes.insert(HeadersToInsert)) { + auto Err = Result.add(Replace); + if (Err) { + consumeError(std::move(Err)); + unsigned NewOffset = Result.getShiftedCodePosition(Replace.getOffset()); + auto Shifted = tooling::Replacement(FileName, NewOffset, 0, + Replace.getReplacementText()); + Result = Result.merge(tooling::Replacements(Shifted)); } } return Result;
diff --git a/lib/Tooling/Inclusions/HeaderIncludes.cpp b/lib/Tooling/Inclusions/HeaderIncludes.cpp index c3bbf6b..285d4a4 100644 --- a/lib/Tooling/Inclusions/HeaderIncludes.cpp +++ b/lib/Tooling/Inclusions/HeaderIncludes.cpp
@@ -14,6 +14,8 @@ #include "clang/Lex/Token.h" #include "clang/Tooling/Core/Replacement.h" #include "clang/Tooling/Inclusions/IncludeStyle.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Error.h" @@ -25,10 +27,12 @@ #include <cassert> #include <climits> #include <functional> +#include <iterator> #include <optional> #include <string> #include <type_traits> #include <utility> +#include <vector> namespace clang { namespace tooling { @@ -512,11 +516,83 @@ return tooling::Replacement(FileName, InsertOffset, 0, NewInclude); } -tooling::Replacements HeaderIncludes::remove(llvm::StringRef IncludeName, - bool IsAngled) const { - assert(IncludeName == trimInclude(IncludeName)); +HeaderIncludes::HeaderToInsert::HeaderToInsert( + llvm::StringRef RawOrSpelledHeader, IncludeDirective Directive, + QuoteStyle QuoteStyle) + : Directive(Directive) { + if (RawOrSpelledHeader.starts_with("<")) { + Header = RawOrSpelledHeader.trim("<>").str(); + this->IsAngled = QuoteStyle != QuoteStyle::QUOTED; + } else if (RawOrSpelledHeader.starts_with("\"")) { + Header = RawOrSpelledHeader.trim("\"").str(); + this->IsAngled = QuoteStyle == QuoteStyle::ANGLED; + } +} + +tooling::Replacements +HeaderIncludes::insert(llvm::ArrayRef<HeaderToInsert> Headers) const { tooling::Replacements Result; - auto Iter = ExistingIncludes.find(IncludeName); + if (Headers.empty()) + return Result; + + std::vector<HeaderToInsert> SortedHeaders = Headers.vec(); + llvm::stable_sort(SortedHeaders, [&](const HeaderToInsert &L, + const HeaderToInsert &R) { + std::string QuotedL = + std::string(llvm::formatv(L.IsAngled ? "<{0}>" : "\"{0}\"", L.Header)); + std::string QuotedR = + std::string(llvm::formatv(R.IsAngled ? "<{0}>" : "\"{0}\"", R.Header)); + int PriorityL = Categories.getIncludePriority( + QuotedL, /*CheckMainHeader=*/!MainIncludeFound); + int PriorityR = Categories.getIncludePriority( + QuotedR, /*CheckMainHeader=*/!MainIncludeFound); + if (PriorityL != PriorityR) + return PriorityL < PriorityR; + if (L.Header != R.Header) + return L.Header < R.Header; + if (L.IsAngled != R.IsAngled) + return L.IsAngled < R.IsAngled; + return L.Directive > R.Directive; + }); + SortedHeaders.erase( + std::unique(SortedHeaders.begin(), SortedHeaders.end(), + [](const HeaderToInsert &L, const HeaderToInsert &R) { + return L.Header == R.Header && L.IsAngled == R.IsAngled; + }), + SortedHeaders.end()); + + struct InsertionInfo { + std::string Text; + unsigned Length = 0; + }; + llvm::DenseMap<unsigned, InsertionInfo> InsertionsByOffset; + + for (const auto &H : SortedHeaders) { + if (auto Insertion = insert(H.Header, H.IsAngled, H.Directive)) { + auto &Info = InsertionsByOffset[Insertion->getOffset()]; + Info.Text += Insertion->getReplacementText(); + if (Insertion->getLength() > 0) { + assert(Info.Length == 0 && "Multiple replacements at same offset?"); + Info.Length = Insertion->getLength(); + } + } + } + + for (const auto &Entry : InsertionsByOffset) { + const auto &Info = Entry.second; + const unsigned Offset = Entry.first; + cantFail(Result.add( + tooling::Replacement(FileName, Offset, Info.Length, Info.Text))); + } + + return Result; +} + +tooling::Replacements HeaderIncludes::remove(llvm::StringRef Header, + bool IsAngled) const { + assert(Header == trimInclude(Header)); + tooling::Replacements Result; + auto Iter = ExistingIncludes.find(Header); if (Iter == ExistingIncludes.end()) return Result; for (const auto &Inc : Iter->second) {
diff --git a/unittests/Format/CleanupTest.cpp b/unittests/Format/CleanupTest.cpp index 06ec780..c36f5ae 100644 --- a/unittests/Format/CleanupTest.cpp +++ b/unittests/Format/CleanupTest.cpp
@@ -464,6 +464,39 @@ EXPECT_EQ(Expected, formatAndApply(Code, Replaces)); } +TEST_F(CleanUpReplacementsTest, InsertMainHeaderAtTopAndSortLLVM) { + std::string Code = "int x;\n"; + std::string Expected = "#include \"fix.h\"\n" + "#include \"a.h\"\n" + "#include \"b.h\"\n" + "#include <vector>\n" + "int x;\n"; + tooling::Replacements Replaces = toReplacements({ + createInsertion("#include \"b.h\""), + createInsertion("#include <vector>"), + createInsertion("#include \"a.h\""), + createInsertion("#include \"fix.h\""), + }); + EXPECT_EQ(Expected, apply(Code, Replaces)); +} + +TEST_F(CleanUpReplacementsTest, InsertMainHeaderAtTopAndSortGoogle) { + std::string Code = "int x;\n"; + std::string Expected = "#include \"fix.h\"\n" + "#include <vector>\n" + "#include \"a.h\"\n" + "#include \"b.h\"\n" + "int x;\n"; + tooling::Replacements Replaces = toReplacements({ + createInsertion("#include \"b.h\""), + createInsertion("#include <vector>"), + createInsertion("#include \"a.h\""), + createInsertion("#include \"fix.h\""), + }); + Style = getGoogleStyle(FormatStyle::LK_Cpp); + EXPECT_EQ(Expected, apply(Code, Replaces)); +} + TEST_F(CleanUpReplacementsTest, NoNewLineAtTheEndOfCodeMultipleInsertions) { std::string Code = "#include <map>"; // FIXME: a better behavior is to only append on newline to Code, but this
diff --git a/unittests/Tooling/HeaderIncludesTest.cpp b/unittests/Tooling/HeaderIncludesTest.cpp index e697538..fdae072 100644 --- a/unittests/Tooling/HeaderIncludesTest.cpp +++ b/unittests/Tooling/HeaderIncludesTest.cpp
@@ -43,6 +43,16 @@ return *Result; } + std::string + insertBatch(llvm::StringRef Code, + llvm::ArrayRef<HeaderIncludes::HeaderToInsert> Headers) { + HeaderIncludes Includes(FileName, Code, Style); + auto Replaces = Includes.insert(Headers); + auto Result = applyAllReplacements(Code, Replaces); + EXPECT_TRUE(static_cast<bool>(Result)); + return *Result; + } + std::string FileName = "fix.cpp"; IncludeStyle Style = format::getLLVMStyle().IncludeStyle; }; @@ -868,6 +878,117 @@ EXPECT_EQ(Expected, insert(Code, "<vector>")); } +TEST_F(HeaderIncludesTest, BatchInsertSortingAndGrouping) { + std::string Code = R"cpp( +void test(); +)cpp"; + std::string Expected = + "\n#include \"a.h\"\n#include \"b.h\"\n#include <foo>\nvoid test();\n"; + + EXPECT_EQ(Expected, + insertBatch(Code, {HeaderIncludes::HeaderToInsert("\"b.h\""), + HeaderIncludes::HeaderToInsert("\"a.h\""), + HeaderIncludes::HeaderToInsert("<foo>")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertMainHeader) { + FileName = "foo.cc"; + Style.IncludeIsMainRegex = "$"; + std::string Code = R"cpp( +void test(); +)cpp"; + std::string Expected = "\n#include \"foo.h\"\n#include \"a.h\"\n#include " + "\"b.h\"\n#include <vector>\nvoid test();\n"; + + EXPECT_EQ(Expected, + insertBatch(Code, {HeaderIncludes::HeaderToInsert("\"b.h\""), + HeaderIncludes::HeaderToInsert("\"foo.h\""), + HeaderIncludes::HeaderToInsert("\"a.h\""), + HeaderIncludes::HeaderToInsert("<vector>")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertDuplicateHeadersSameQuotes) { + std::string Code = R"cpp( +void test(); +)cpp"; + std::string Expected = "\n#include \"a.h\"\n#include <foo>\nvoid test();\n"; + + EXPECT_EQ(Expected, + insertBatch(Code, {HeaderIncludes::HeaderToInsert("\"a.h\""), + HeaderIncludes::HeaderToInsert("\"a.h\""), + HeaderIncludes::HeaderToInsert("<foo>")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertDuplicateHeadersMixedQuotes) { + std::string Code = R"cpp( +void test(); +)cpp"; + std::string Expected = "\n#include \"a.h\"\n#include <a.h>\nvoid test();\n"; + + EXPECT_EQ(Expected, + insertBatch(Code, {HeaderIncludes::HeaderToInsert("\"a.h\""), + HeaderIncludes::HeaderToInsert("<a.h>")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertDuplicateHeadersMixedDirectives) { + std::string Code = R"cpp( +void test(); +)cpp"; + std::string Expected = "\n#import \"a.h\"\n#include <foo>\nvoid test();\n"; + + EXPECT_EQ( + Expected, + insertBatch( + Code, + {HeaderIncludes::HeaderToInsert("\"a.h\"", IncludeDirective::Include), + HeaderIncludes::HeaderToInsert("\"a.h\"", IncludeDirective::Import), + HeaderIncludes::HeaderToInsert("<foo>")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertIntoExistingIncludes) { + std::string Code = R"cpp(#include "b.h" + +void test(); +)cpp"; + std::string Expected = R"cpp(#include "a.h" +#include "b.h" +#include "c.h" +#include <vector> + +void test(); +)cpp"; + + EXPECT_EQ(Expected, + insertBatch(Code, {HeaderIncludes::HeaderToInsert("\"c.h\""), + HeaderIncludes::HeaderToInsert("\"a.h\""), + HeaderIncludes::HeaderToInsert("<vector>")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertEmptyFile) { + std::string Code = ""; + std::string Expected = "#include \"a.h\"\n#include <foo>\n"; + + EXPECT_EQ(Expected, + insertBatch(Code, {HeaderIncludes::HeaderToInsert("<foo>"), + HeaderIncludes::HeaderToInsert("\"a.h\"")})); +} + +TEST_F(HeaderIncludesTest, BatchInsertImportsAlongsideIncludes) { + std::string Code = R"objc( +void test(); +)objc"; + std::string Expected = + "\n#import \"a.h\"\n#include \"b.h\"\n#import <foo>\nvoid test();\n"; + + EXPECT_EQ( + Expected, + insertBatch( + Code, + {HeaderIncludes::HeaderToInsert("\"b.h\"", IncludeDirective::Include), + HeaderIncludes::HeaderToInsert("\"a.h\"", IncludeDirective::Import), + HeaderIncludes::HeaderToInsert("<foo>", IncludeDirective::Import)})); +} + } // namespace } // namespace tooling } // namespace clang