1 //===-------- IncludeInserter.cpp - clang-tidy ----------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "IncludeInserter.h"
10 #include "clang/Lex/PPCallbacks.h"
11 #include "clang/Lex/Preprocessor.h"
12 #include "clang/Lex/Token.h"
13
14 namespace clang {
15 namespace tidy {
16 namespace utils {
17
18 class IncludeInserterCallback : public PPCallbacks {
19 public:
IncludeInserterCallback(IncludeInserter * Inserter)20 explicit IncludeInserterCallback(IncludeInserter *Inserter)
21 : Inserter(Inserter) {}
22 // Implements PPCallbacks::InclusionDerective(). Records the names and source
23 // locations of the inclusions in the main source file being processed.
InclusionDirective(SourceLocation HashLocation,const Token & IncludeToken,StringRef FileNameRef,bool IsAngled,CharSourceRange FileNameRange,const FileEntry *,StringRef,StringRef,const Module *,SrcMgr::CharacteristicKind)24 void InclusionDirective(SourceLocation HashLocation,
25 const Token &IncludeToken, StringRef FileNameRef,
26 bool IsAngled, CharSourceRange FileNameRange,
27 const FileEntry * /*IncludedFile*/,
28 StringRef /*SearchPath*/, StringRef /*RelativePath*/,
29 const Module * /*ImportedModule*/,
30 SrcMgr::CharacteristicKind /*FileType*/) override {
31 Inserter->addInclude(FileNameRef, IsAngled, HashLocation,
32 IncludeToken.getEndLoc());
33 }
34
35 private:
36 IncludeInserter *Inserter;
37 };
38
IncludeInserter(IncludeSorter::IncludeStyle Style)39 IncludeInserter::IncludeInserter(IncludeSorter::IncludeStyle Style)
40 : Style(Style) {}
41
registerPreprocessor(Preprocessor * PP)42 void IncludeInserter::registerPreprocessor(Preprocessor *PP) {
43 assert(PP && "PP shouldn't be null");
44 SourceMgr = &PP->getSourceManager();
45
46 // If this gets registered multiple times, clear the maps
47 if (!IncludeSorterByFile.empty())
48 IncludeSorterByFile.clear();
49 if (!InsertedHeaders.empty())
50 InsertedHeaders.clear();
51 PP->addPPCallbacks(std::make_unique<IncludeInserterCallback>(this));
52 }
53
getOrCreate(FileID FileID)54 IncludeSorter &IncludeInserter::getOrCreate(FileID FileID) {
55 assert(SourceMgr && "SourceMgr shouldn't be null; did you remember to call "
56 "registerPreprocessor()?");
57 // std::unique_ptr is cheap to construct, so force a construction now to save
58 // the lookup needed if we were to insert into the map.
59 std::unique_ptr<IncludeSorter> &Entry = IncludeSorterByFile[FileID];
60 if (!Entry) {
61 // If it wasn't found, Entry will be default constructed to nullptr.
62 Entry = std::make_unique<IncludeSorter>(
63 SourceMgr, FileID,
64 SourceMgr->getFilename(SourceMgr->getLocForStartOfFile(FileID)), Style);
65 }
66 return *Entry;
67 }
68
69 llvm::Optional<FixItHint>
createIncludeInsertion(FileID FileID,llvm::StringRef Header)70 IncludeInserter::createIncludeInsertion(FileID FileID, llvm::StringRef Header) {
71 bool IsAngled = Header.consume_front("<");
72 if (IsAngled != Header.consume_back(">"))
73 return llvm::None;
74 // We assume the same Header will never be included both angled and not
75 // angled.
76 if (!InsertedHeaders[FileID].insert(Header).second)
77 return llvm::None;
78
79 return getOrCreate(FileID).CreateIncludeInsertion(Header, IsAngled);
80 }
81
82 llvm::Optional<FixItHint>
createMainFileIncludeInsertion(StringRef Header)83 IncludeInserter::createMainFileIncludeInsertion(StringRef Header) {
84 assert(SourceMgr && "SourceMgr shouldn't be null; did you remember to call "
85 "registerPreprocessor()?");
86 return createIncludeInsertion(SourceMgr->getMainFileID(), Header);
87 }
88
addInclude(StringRef FileName,bool IsAngled,SourceLocation HashLocation,SourceLocation EndLocation)89 void IncludeInserter::addInclude(StringRef FileName, bool IsAngled,
90 SourceLocation HashLocation,
91 SourceLocation EndLocation) {
92 assert(SourceMgr && "SourceMgr shouldn't be null; did you remember to call "
93 "registerPreprocessor()?");
94 FileID FileID = SourceMgr->getFileID(HashLocation);
95 getOrCreate(FileID).AddInclude(FileName, IsAngled, HashLocation, EndLocation);
96 }
97
98 } // namespace utils
99 } // namespace tidy
100 } // namespace clang
101