Index: cfe/trunk/include/clang/Tooling/Refactoring/Transformer.h =================================================================== --- cfe/trunk/include/clang/Tooling/Refactoring/Transformer.h +++ cfe/trunk/include/clang/Tooling/Refactoring/Transformer.h @@ -86,6 +86,12 @@ TextGenerator Note; }; +/// Format of the path in an include directive -- angle brackets or quotes. +enum class IncludeFormat { + Quoted, + Angled, +}; + /// Description of a source-code transformation. // // A *rewrite rule* describes a transformation of source code. A simple rule @@ -114,6 +120,10 @@ ast_matchers::internal::DynTypedMatcher Matcher; SmallVector Edits; TextGenerator Explanation; + // Include paths to add to the file affected by this case. These are + // bundled with the `Case`, rather than the `RewriteRule`, because each case + // might have different associated changes to the includes. + std::vector> AddedIncludes; }; // We expect RewriteRules will most commonly include only one case. SmallVector Cases; @@ -137,6 +147,19 @@ return makeRule(std::move(M), std::move(Edits), std::move(Explanation)); } +/// For every case in Rule, adds an include directive for the given header. The +/// common use is assumed to be a rule with only one case. For example, to +/// replace a function call and add headers corresponding to the new code, one +/// could write: +/// \code +/// auto R = makeRule(callExpr(callee(functionDecl(hasName("foo")))), +/// change(text("bar()"))); +/// AddInclude(R, "path/to/bar_header.h"); +/// AddInclude(R, "vector", IncludeFormat::Angled); +/// \endcode +void addInclude(RewriteRule &Rule, llvm::StringRef Header, + IncludeFormat Format = IncludeFormat::Quoted); + /// Applies the first rule whose pattern matches; other rules are ignored. /// /// N.B. All of the rules must use the same kind of matcher (that is, share a Index: cfe/trunk/lib/Tooling/Refactoring/Transformer.cpp =================================================================== --- cfe/trunk/lib/Tooling/Refactoring/Transformer.cpp +++ cfe/trunk/lib/Tooling/Refactoring/Transformer.cpp @@ -98,8 +98,14 @@ RewriteRule tooling::makeRule(DynTypedMatcher M, SmallVector Edits, TextGenerator Explanation) { - return RewriteRule{{RewriteRule::Case{std::move(M), std::move(Edits), - std::move(Explanation)}}}; + return RewriteRule{{RewriteRule::Case{ + std::move(M), std::move(Edits), std::move(Explanation), {}}}}; +} + +void tooling::addInclude(RewriteRule &Rule, StringRef Header, + IncludeFormat Format) { + for (auto &Case : Rule.Cases) + Case.AddedIncludes.emplace_back(Header.str(), Format); } // Determines whether A is a base type of B in the class hierarchy, including @@ -217,8 +223,8 @@ Root->second.getSourceRange().getBegin()); assert(RootLoc.isValid() && "Invalid location for Root node of match."); - auto Transformations = tooling::detail::translateEdits( - Result, tooling::detail::findSelectedCase(Result, Rule).Edits); + RewriteRule::Case Case = tooling::detail::findSelectedCase(Result, Rule); + auto Transformations = tooling::detail::translateEdits(Result, Case.Edits); if (!Transformations) { Consumer(Transformations.takeError()); return; @@ -241,5 +247,17 @@ } } + for (const auto &I : Case.AddedIncludes) { + auto &Header = I.first; + switch (I.second) { + case IncludeFormat::Quoted: + AC.addHeader(Header); + break; + case IncludeFormat::Angled: + AC.addHeader((llvm::Twine("<") + Header + ">").str()); + break; + } + } + Consumer(std::move(AC)); } Index: cfe/trunk/unittests/Tooling/TransformerTest.cpp =================================================================== --- cfe/trunk/unittests/Tooling/TransformerTest.cpp +++ cfe/trunk/unittests/Tooling/TransformerTest.cpp @@ -198,6 +198,42 @@ testRule(std::move(Rule), Input, Expected); } +TEST_F(TransformerTest, AddIncludeQuoted) { + RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))), + change(text("other()"))); + addInclude(Rule, "clang/OtherLib.h"); + + 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"; + + testRule(Rule, Input, Expected); +} + +TEST_F(TransformerTest, AddIncludeAngled) { + RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))), + change(text("other()"))); + addInclude(Rule, "clang/OtherLib.h", IncludeFormat::Angled); + + 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"; + + testRule(Rule, Input, Expected); +} + TEST_F(TransformerTest, NodePartNameNamedDecl) { StringRef Fun = "fun"; RewriteRule Rule = makeRule(functionDecl(hasName("bad")).bind(Fun),