diff --git a/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.cpp b/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.cpp index 0f7771cd5648..4ff976de14be 100644 --- a/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.cpp +++ b/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.cpp @@ -46,6 +46,19 @@ TransformerClangTidyCheck::TransformerClangTidyCheck(RewriteRule R, " explicitly provide an empty explanation if none is desired"); } +void TransformerClangTidyCheck::registerPPCallbacks( + const SourceManager &SM, Preprocessor *PP, Preprocessor *ModuleExpanderPP) { + // Only allocate and register the IncludeInsert when some `Case` will add + // includes. + if (Rule && llvm::any_of(Rule->Cases, [](const RewriteRule::Case &C) { + return !C.AddedIncludes.empty(); + })) { + Inserter = llvm::make_unique( + SM, getLangOpts(), utils::IncludeSorter::IS_LLVM); + PP->addPPCallbacks(Inserter->CreatePPCallbacks()); + } +} + void TransformerClangTidyCheck::registerMatchers( ast_matchers::MatchFinder *Finder) { if (Rule) @@ -89,6 +102,15 @@ void TransformerClangTidyCheck::check( for (const auto &T : *Transformations) { Diag << FixItHint::CreateReplacement(T.Range, T.Replacement); } + + for (const auto &I : Case.AddedIncludes) { + auto &Header = I.first; + if (Optional Fix = Inserter->CreateIncludeInsertion( + Result.SourceManager->getMainFileID(), Header, + /*IsAngled=*/I.second == tooling::IncludeFormat::Angled)) { + Diag << *Fix; + } + } } } // namespace utils diff --git a/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.h b/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.h index 0cd386f5b07d..9a40bb7397b4 100644 --- a/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.h +++ b/clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.h @@ -10,7 +10,9 @@ #define LLVM_CLANG_TOOLS_EXTRA_CLANG_TIDY_TRANSFORMER_CLANG_TIDY_CHECK_H #include "../ClangTidy.h" +#include "../utils/IncludeInserter.h" #include "clang/ASTMatchers/ASTMatchFinder.h" +#include "clang/Frontend/CompilerInstance.h" #include "clang/Tooling/Refactoring/Transformer.h" #include #include @@ -52,11 +54,14 @@ public: TransformerClangTidyCheck(tooling::RewriteRule R, StringRef Name, ClangTidyContext *Context); + void registerPPCallbacks(const SourceManager &SM, Preprocessor *PP, + Preprocessor *ModuleExpanderPP) override; void registerMatchers(ast_matchers::MatchFinder *Finder) final; void check(const ast_matchers::MatchFinder::MatchResult &Result) final; private: Optional Rule; + std::unique_ptr Inserter; }; } // namespace utils diff --git a/clang-tools-extra/unittests/clang-tidy/TransformerClangTidyCheckTest.cpp b/clang-tools-extra/unittests/clang-tidy/TransformerClangTidyCheckTest.cpp index 6e42be25d5f6..9a3457fe252c 100644 --- a/clang-tools-extra/unittests/clang-tidy/TransformerClangTidyCheckTest.cpp +++ b/clang-tools-extra/unittests/clang-tidy/TransformerClangTidyCheckTest.cpp @@ -19,6 +19,7 @@ namespace tidy { namespace utils { namespace { using tooling::change; +using tooling::IncludeFormat; using tooling::RewriteRule; using tooling::text; using tooling::stencil::cat; @@ -121,6 +122,54 @@ TEST(TransformerClangTidyCheckTest, DisableByConfig) { Input, nullptr, "input.cc", None, Options)); } +RewriteRule replaceCall(IncludeFormat Format) { + using namespace ::clang::ast_matchers; + RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))), + change(text("other()")), text("no message")); + addInclude(Rule, "clang/OtherLib.h", Format); + return Rule; +} + +template +class IncludeCheck : public TransformerClangTidyCheck { +public: + IncludeCheck(StringRef Name, ClangTidyContext *Context) + : TransformerClangTidyCheck(replaceCall(Format), Name, Context) {} +}; + +TEST(TransformerClangTidyCheckTest, AddIncludeQuoted) { + + std::string Input = R"cc( + int f(int x); + int h(int x) { return f(x); } + )cc"; + std::string Expected = R"cc(#include "clang/OtherLib.h" + + + int f(int x); + int h(int x) { return other(); } + )cc"; + + EXPECT_EQ(Expected, + test::runCheckOnCode>(Input)); +} + +TEST(TransformerClangTidyCheckTest, AddIncludeAngled) { + std::string Input = R"cc( + int f(int x); + int h(int x) { return f(x); } + )cc"; + std::string Expected = R"cc(#include + + + int f(int x); + int h(int x) { return other(); } + )cc"; + + EXPECT_EQ(Expected, + test::runCheckOnCode>(Input)); +} + } // namespace } // namespace utils } // namespace tidy