2 files changed, 54 insertions, 8 deletions

diff --git a/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.cpp b/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.cpp
index ff696597357..fd2b74cba18 100644
--- a/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.cpp
+++ b/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.cpp
@@ -8,7 +8,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "MakeSharedCheck.h"
+#include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/Lexer.h"
+#include "clang/Lex/Preprocessor.h"
 
 using namespace clang::ast_matchers;
 
@@ -17,6 +19,9 @@ namespace tidy {
 namespace modernize {
 
 namespace {
+
+constexpr char StdMemoryHeader[] = "memory";
+
 std::string GetNewExprName(const CXXNewExpr *NewExpr,
                            const SourceManager &SM,
                            const LangOptions &Lang) {
@@ -29,6 +34,7 @@ std::string GetNewExprName(const CXXNewExpr *NewExpr,
   }
   return WrittenName.str();
 }
+
 } // namespace
 
 const char MakeSmartPtrCheck::PointerType[] = "pointerType";
@@ -37,9 +43,28 @@ const char MakeSmartPtrCheck::ResetCall[] = "resetCall";
 const char MakeSmartPtrCheck::NewExpression[] = "newExpression";
 
 MakeSmartPtrCheck::MakeSmartPtrCheck(StringRef Name, ClangTidyContext *Context,
-                                     std::string makeSmartPtrFunctionName)
+                                     StringRef MakeSmartPtrFunctionName)
     : ClangTidyCheck(Name, Context),
-      makeSmartPtrFunctionName(std::move(makeSmartPtrFunctionName)) {}
+      IncludeStyle(utils::IncludeSorter::parseIncludeStyle(
+          Options.get("IncludeStyle", "llvm"))),
+      MakeSmartPtrFunctionHeader(
+          Options.get("MakeSmartPtrFunctionHeader", StdMemoryHeader)),
+      MakeSmartPtrFunctionName(
+          Options.get("MakeSmartPtrFunction", MakeSmartPtrFunctionName)) {}
+
+void MakeSmartPtrCheck::storeOptions(ClangTidyOptions::OptionMap &Opts) {
+  Options.store(Opts, "IncludeStyle", IncludeStyle);
+  Options.store(Opts, "MakeSmartPtrFunctionHeader", MakeSmartPtrFunctionHeader);
+  Options.store(Opts, "MakeSmartPtrFunction", MakeSmartPtrFunctionName);
+}
+
+void MakeSmartPtrCheck::registerPPCallbacks(CompilerInstance &Compiler) {
+  if (getLangOpts().CPlusPlus11) {
+    Inserter.reset(new utils::IncludeInserter(
+        Compiler.getSourceManager(), Compiler.getLangOpts(), IncludeStyle));
+    Compiler.getPreprocessor().addPPCallbacks(Inserter->CreatePPCallbacks());
+  }
+}
 
 void MakeSmartPtrCheck::registerMatchers(ast_matchers::MatchFinder *Finder) {
   if (!getLangOpts().CPlusPlus11)
@@ -107,7 +132,7 @@ void MakeSmartPtrCheck::checkConstruct(SourceManager &SM,
     return;
 
   auto Diag = diag(ConstructCallStart, "use %0 instead")
-              << makeSmartPtrFunctionName;
+              << MakeSmartPtrFunctionName;
 
   // Find the location of the template's left angle.
   size_t LAngle = ExprStr.find("<");
@@ -125,7 +150,7 @@ void MakeSmartPtrCheck::checkConstruct(SourceManager &SM,
 
   Diag << FixItHint::CreateReplacement(
       CharSourceRange::getCharRange(ConstructCallStart, ConstructCallEnd),
-      makeSmartPtrFunctionName);
+      MakeSmartPtrFunctionName);
 
   // If the smart_ptr is built with brace enclosed direct initialization, use
   // parenthesis instead.
@@ -142,6 +167,7 @@ void MakeSmartPtrCheck::checkConstruct(SourceManager &SM,
   }
 
   replaceNew(Diag, New, SM);
+  insertHeader(Diag, SM.getFileID(ConstructCallStart));
 }
 
 void MakeSmartPtrCheck::checkReset(SourceManager &SM,
@@ -155,11 +181,11 @@ void MakeSmartPtrCheck::checkReset(SourceManager &SM,
       Lexer::getLocForEndOfToken(Expr->getLocEnd(), 0, SM, getLangOpts());
 
   auto Diag = diag(ResetCallStart, "use %0 instead")
-              << makeSmartPtrFunctionName;
+              << MakeSmartPtrFunctionName;
 
   Diag << FixItHint::CreateReplacement(
       CharSourceRange::getCharRange(OperatorLoc, ExprEnd),
-      (llvm::Twine(" = ") + makeSmartPtrFunctionName + "<" +
+      (llvm::Twine(" = ") + MakeSmartPtrFunctionName + "<" +
        GetNewExprName(New, SM, getLangOpts()) + ">")
           .str());
 
@@ -167,6 +193,7 @@ void MakeSmartPtrCheck::checkReset(SourceManager &SM,
     Diag << FixItHint::CreateInsertion(ExprStart, "*");
 
   replaceNew(Diag, New, SM);
+  insertHeader(Diag, SM.getFileID(OperatorLoc));
 }
 
 void MakeSmartPtrCheck::replaceNew(DiagnosticBuilder &Diag,
@@ -243,6 +270,17 @@ void MakeSmartPtrCheck::replaceNew(DiagnosticBuilder &Diag,
   }
 }
 
+void MakeSmartPtrCheck::insertHeader(DiagnosticBuilder &Diag, FileID FD) {
+  if (MakeSmartPtrFunctionHeader.empty()) {
+    return;
+  }
+  if (auto IncludeFixit = Inserter->CreateIncludeInsertion(
+          FD, MakeSmartPtrFunctionHeader,
+          /*IsAngled=*/MakeSmartPtrFunctionHeader == StdMemoryHeader)) {
+    Diag << *IncludeFixit;
+  }
+}
+
 } // namespace modernize
 } // namespace tidy
 } // namespace clang
diff --git a/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.h b/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.h
index 0d6d0c9d54a..f01550e28fd 100644
--- a/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.h
+++ b/clang-tools-extra/clang-tidy/modernize/MakeSmartPtrCheck.h
@@ -11,6 +11,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANG_TIDY_MODERNIZE_MAKE_SMART_PTR_H
 
 #include "../ClangTidy.h"
+#include "../utils/IncludeInserter.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchersInternal.h"
 #include "llvm/ADT/StringRef.h"
@@ -24,9 +25,11 @@ namespace modernize {
 class MakeSmartPtrCheck : public ClangTidyCheck {
 public:
   MakeSmartPtrCheck(StringRef Name, ClangTidyContext *Context,
-                    std::string makeSmartPtrFunctionName);
+                    StringRef MakeSmartPtrFunctionName);
   void registerMatchers(ast_matchers::MatchFinder *Finder) final;
+  void registerPPCallbacks(clang::CompilerInstance &Compiler) override;
   void check(const ast_matchers::MatchFinder::MatchResult &Result) final;
+  void storeOptions(ClangTidyOptions::OptionMap &Opts) override;
 
 protected:
   using SmartPtrTypeMatcher = ast_matchers::internal::BindableMatcher<QualType>;
@@ -43,14 +46,19 @@ protected:
   static const char NewExpression[];
 
 private:
-  std::string makeSmartPtrFunctionName;
+  std::unique_ptr<utils::IncludeInserter> Inserter;
+  const utils::IncludeSorter::IncludeStyle IncludeStyle;
+  const std::string MakeSmartPtrFunctionHeader;
+  const std::string MakeSmartPtrFunctionName;
 
   void checkConstruct(SourceManager &SM, const CXXConstructExpr *Construct,
                       const QualType *Type, const CXXNewExpr *New);
   void checkReset(SourceManager &SM, const CXXMemberCallExpr *Member,
                   const CXXNewExpr *New);
+
   void replaceNew(DiagnosticBuilder &Diag, const CXXNewExpr *New,
                   SourceManager &SM);
+  void insertHeader(DiagnosticBuilder &Diag, FileID FD);
 };
 
 } // namespace modernize