xref: /openbsd-src/gnu/llvm/clang/lib/Tooling/Inclusions/HeaderIncludes.cpp (revision 12c855180aad702bbcca06e0398d774beeafb155)
1e5dd7070Spatrick //===--- HeaderIncludes.cpp - Insert/Delete #includes --*- C++ -*----------===//
2e5dd7070Spatrick //
3e5dd7070Spatrick // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4e5dd7070Spatrick // See https://llvm.org/LICENSE.txt for license information.
5e5dd7070Spatrick // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6e5dd7070Spatrick //
7e5dd7070Spatrick //===----------------------------------------------------------------------===//
8e5dd7070Spatrick 
9e5dd7070Spatrick #include "clang/Tooling/Inclusions/HeaderIncludes.h"
10ec727ea7Spatrick #include "clang/Basic/FileManager.h"
11e5dd7070Spatrick #include "clang/Basic/SourceManager.h"
12e5dd7070Spatrick #include "clang/Lex/Lexer.h"
13e5dd7070Spatrick #include "llvm/Support/FormatVariadic.h"
14a9ac8606Spatrick #include "llvm/Support/Path.h"
15*12c85518Srobert #include <optional>
16e5dd7070Spatrick 
17e5dd7070Spatrick namespace clang {
18e5dd7070Spatrick namespace tooling {
19e5dd7070Spatrick namespace {
20e5dd7070Spatrick 
createLangOpts()21e5dd7070Spatrick LangOptions createLangOpts() {
22e5dd7070Spatrick   LangOptions LangOpts;
23e5dd7070Spatrick   LangOpts.CPlusPlus = 1;
24e5dd7070Spatrick   LangOpts.CPlusPlus11 = 1;
25e5dd7070Spatrick   LangOpts.CPlusPlus14 = 1;
26e5dd7070Spatrick   LangOpts.LineComment = 1;
27e5dd7070Spatrick   LangOpts.CXXOperatorNames = 1;
28e5dd7070Spatrick   LangOpts.Bool = 1;
29e5dd7070Spatrick   LangOpts.ObjC = 1;
30e5dd7070Spatrick   LangOpts.MicrosoftExt = 1;    // To get kw___try, kw___finally.
31e5dd7070Spatrick   LangOpts.DeclSpecKeyword = 1; // To get __declspec.
32e5dd7070Spatrick   LangOpts.WChar = 1;           // To get wchar_t
33e5dd7070Spatrick   return LangOpts;
34e5dd7070Spatrick }
35e5dd7070Spatrick 
36e5dd7070Spatrick // Returns the offset after skipping a sequence of tokens, matched by \p
37e5dd7070Spatrick // GetOffsetAfterSequence, from the start of the code.
38e5dd7070Spatrick // \p GetOffsetAfterSequence should be a function that matches a sequence of
39e5dd7070Spatrick // tokens and returns an offset after the sequence.
getOffsetAfterTokenSequence(StringRef FileName,StringRef Code,const IncludeStyle & Style,llvm::function_ref<unsigned (const SourceManager &,Lexer &,Token &)> GetOffsetAfterSequence)40e5dd7070Spatrick unsigned getOffsetAfterTokenSequence(
41e5dd7070Spatrick     StringRef FileName, StringRef Code, const IncludeStyle &Style,
42e5dd7070Spatrick     llvm::function_ref<unsigned(const SourceManager &, Lexer &, Token &)>
43e5dd7070Spatrick         GetOffsetAfterSequence) {
44e5dd7070Spatrick   SourceManagerForFile VirtualSM(FileName, Code);
45e5dd7070Spatrick   SourceManager &SM = VirtualSM.get();
46*12c85518Srobert   LangOptions LangOpts = createLangOpts();
47a9ac8606Spatrick   Lexer Lex(SM.getMainFileID(), SM.getBufferOrFake(SM.getMainFileID()), SM,
48*12c85518Srobert             LangOpts);
49e5dd7070Spatrick   Token Tok;
50e5dd7070Spatrick   // Get the first token.
51e5dd7070Spatrick   Lex.LexFromRawLexer(Tok);
52e5dd7070Spatrick   return GetOffsetAfterSequence(SM, Lex, Tok);
53e5dd7070Spatrick }
54e5dd7070Spatrick 
55e5dd7070Spatrick // Check if a sequence of tokens is like "#<Name> <raw_identifier>". If it is,
56e5dd7070Spatrick // \p Tok will be the token after this directive; otherwise, it can be any token
57e5dd7070Spatrick // after the given \p Tok (including \p Tok). If \p RawIDName is provided, the
58e5dd7070Spatrick // (second) raw_identifier name is checked.
checkAndConsumeDirectiveWithName(Lexer & Lex,StringRef Name,Token & Tok,std::optional<StringRef> RawIDName=std::nullopt)59e5dd7070Spatrick bool checkAndConsumeDirectiveWithName(
60e5dd7070Spatrick     Lexer &Lex, StringRef Name, Token &Tok,
61*12c85518Srobert     std::optional<StringRef> RawIDName = std::nullopt) {
62e5dd7070Spatrick   bool Matched = Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
63e5dd7070Spatrick                  Tok.is(tok::raw_identifier) &&
64e5dd7070Spatrick                  Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) &&
65e5dd7070Spatrick                  Tok.is(tok::raw_identifier) &&
66e5dd7070Spatrick                  (!RawIDName || Tok.getRawIdentifier() == *RawIDName);
67e5dd7070Spatrick   if (Matched)
68e5dd7070Spatrick     Lex.LexFromRawLexer(Tok);
69e5dd7070Spatrick   return Matched;
70e5dd7070Spatrick }
71e5dd7070Spatrick 
skipComments(Lexer & Lex,Token & Tok)72e5dd7070Spatrick void skipComments(Lexer &Lex, Token &Tok) {
73e5dd7070Spatrick   while (Tok.is(tok::comment))
74e5dd7070Spatrick     if (Lex.LexFromRawLexer(Tok))
75e5dd7070Spatrick       return;
76e5dd7070Spatrick }
77e5dd7070Spatrick 
78e5dd7070Spatrick // Returns the offset after header guard directives and any comments
79e5dd7070Spatrick // before/after header guards (e.g. #ifndef/#define pair, #pragma once). If no
80e5dd7070Spatrick // header guard is present in the code, this will return the offset after
81e5dd7070Spatrick // skipping all comments from the start of the code.
getOffsetAfterHeaderGuardsAndComments(StringRef FileName,StringRef Code,const IncludeStyle & Style)82e5dd7070Spatrick unsigned getOffsetAfterHeaderGuardsAndComments(StringRef FileName,
83e5dd7070Spatrick                                                StringRef Code,
84e5dd7070Spatrick                                                const IncludeStyle &Style) {
85e5dd7070Spatrick   // \p Consume returns location after header guard or 0 if no header guard is
86e5dd7070Spatrick   // found.
87e5dd7070Spatrick   auto ConsumeHeaderGuardAndComment =
88e5dd7070Spatrick       [&](std::function<unsigned(const SourceManager &SM, Lexer &Lex,
89e5dd7070Spatrick                                  Token Tok)>
90e5dd7070Spatrick               Consume) {
91e5dd7070Spatrick         return getOffsetAfterTokenSequence(
92e5dd7070Spatrick             FileName, Code, Style,
93e5dd7070Spatrick             [&Consume](const SourceManager &SM, Lexer &Lex, Token Tok) {
94e5dd7070Spatrick               skipComments(Lex, Tok);
95e5dd7070Spatrick               unsigned InitialOffset = SM.getFileOffset(Tok.getLocation());
96e5dd7070Spatrick               return std::max(InitialOffset, Consume(SM, Lex, Tok));
97e5dd7070Spatrick             });
98e5dd7070Spatrick       };
99e5dd7070Spatrick   return std::max(
100e5dd7070Spatrick       // #ifndef/#define
101e5dd7070Spatrick       ConsumeHeaderGuardAndComment(
102e5dd7070Spatrick           [](const SourceManager &SM, Lexer &Lex, Token Tok) -> unsigned {
103e5dd7070Spatrick             if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) {
104e5dd7070Spatrick               skipComments(Lex, Tok);
105a9ac8606Spatrick               if (checkAndConsumeDirectiveWithName(Lex, "define", Tok) &&
106a9ac8606Spatrick                   Tok.isAtStartOfLine())
107e5dd7070Spatrick                 return SM.getFileOffset(Tok.getLocation());
108e5dd7070Spatrick             }
109e5dd7070Spatrick             return 0;
110e5dd7070Spatrick           }),
111e5dd7070Spatrick       // #pragma once
112e5dd7070Spatrick       ConsumeHeaderGuardAndComment(
113e5dd7070Spatrick           [](const SourceManager &SM, Lexer &Lex, Token Tok) -> unsigned {
114e5dd7070Spatrick             if (checkAndConsumeDirectiveWithName(Lex, "pragma", Tok,
115e5dd7070Spatrick                                                  StringRef("once")))
116e5dd7070Spatrick               return SM.getFileOffset(Tok.getLocation());
117e5dd7070Spatrick             return 0;
118e5dd7070Spatrick           }));
119e5dd7070Spatrick }
120e5dd7070Spatrick 
121e5dd7070Spatrick // Check if a sequence of tokens is like
122e5dd7070Spatrick //    "#include ("header.h" | <header.h>)".
123e5dd7070Spatrick // If it is, \p Tok will be the token after this directive; otherwise, it can be
124e5dd7070Spatrick // any token after the given \p Tok (including \p Tok).
checkAndConsumeInclusiveDirective(Lexer & Lex,Token & Tok)125e5dd7070Spatrick bool checkAndConsumeInclusiveDirective(Lexer &Lex, Token &Tok) {
126e5dd7070Spatrick   auto Matched = [&]() {
127e5dd7070Spatrick     Lex.LexFromRawLexer(Tok);
128e5dd7070Spatrick     return true;
129e5dd7070Spatrick   };
130e5dd7070Spatrick   if (Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
131e5dd7070Spatrick       Tok.is(tok::raw_identifier) && Tok.getRawIdentifier() == "include") {
132e5dd7070Spatrick     if (Lex.LexFromRawLexer(Tok))
133e5dd7070Spatrick       return false;
134e5dd7070Spatrick     if (Tok.is(tok::string_literal))
135e5dd7070Spatrick       return Matched();
136e5dd7070Spatrick     if (Tok.is(tok::less)) {
137e5dd7070Spatrick       while (!Lex.LexFromRawLexer(Tok) && Tok.isNot(tok::greater)) {
138e5dd7070Spatrick       }
139e5dd7070Spatrick       if (Tok.is(tok::greater))
140e5dd7070Spatrick         return Matched();
141e5dd7070Spatrick     }
142e5dd7070Spatrick   }
143e5dd7070Spatrick   return false;
144e5dd7070Spatrick }
145e5dd7070Spatrick 
146e5dd7070Spatrick // Returns the offset of the last #include directive after which a new
147e5dd7070Spatrick // #include can be inserted. This ignores #include's after the #include block(s)
148e5dd7070Spatrick // in the beginning of a file to avoid inserting headers into code sections
149e5dd7070Spatrick // where new #include's should not be added by default.
150e5dd7070Spatrick // These code sections include:
151e5dd7070Spatrick //      - raw string literals (containing #include).
152e5dd7070Spatrick //      - #if blocks.
153e5dd7070Spatrick //      - Special #include's among declarations (e.g. functions).
154e5dd7070Spatrick //
155e5dd7070Spatrick // If no #include after which a new #include can be inserted, this returns the
156e5dd7070Spatrick // offset after skipping all comments from the start of the code.
157e5dd7070Spatrick // Inserting after an #include is not allowed if it comes after code that is not
158e5dd7070Spatrick // #include (e.g. pre-processing directive that is not #include, declarations).
getMaxHeaderInsertionOffset(StringRef FileName,StringRef Code,const IncludeStyle & Style)159e5dd7070Spatrick unsigned getMaxHeaderInsertionOffset(StringRef FileName, StringRef Code,
160e5dd7070Spatrick                                      const IncludeStyle &Style) {
161e5dd7070Spatrick   return getOffsetAfterTokenSequence(
162e5dd7070Spatrick       FileName, Code, Style,
163e5dd7070Spatrick       [](const SourceManager &SM, Lexer &Lex, Token Tok) {
164e5dd7070Spatrick         skipComments(Lex, Tok);
165e5dd7070Spatrick         unsigned MaxOffset = SM.getFileOffset(Tok.getLocation());
166e5dd7070Spatrick         while (checkAndConsumeInclusiveDirective(Lex, Tok))
167e5dd7070Spatrick           MaxOffset = SM.getFileOffset(Tok.getLocation());
168e5dd7070Spatrick         return MaxOffset;
169e5dd7070Spatrick       });
170e5dd7070Spatrick }
171e5dd7070Spatrick 
trimInclude(StringRef IncludeName)172e5dd7070Spatrick inline StringRef trimInclude(StringRef IncludeName) {
173e5dd7070Spatrick   return IncludeName.trim("\"<>");
174e5dd7070Spatrick }
175e5dd7070Spatrick 
176e5dd7070Spatrick const char IncludeRegexPattern[] =
177e5dd7070Spatrick     R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))";
178e5dd7070Spatrick 
179a9ac8606Spatrick // The filename of Path excluding extension.
180a9ac8606Spatrick // Used to match implementation with headers, this differs from sys::path::stem:
181a9ac8606Spatrick //  - in names with multiple dots (foo.cu.cc) it terminates at the *first*
182a9ac8606Spatrick //  - an empty stem is never returned: /foo/.bar.x => .bar
183a9ac8606Spatrick //  - we don't bother to handle . and .. specially
matchingStem(llvm::StringRef Path)184a9ac8606Spatrick StringRef matchingStem(llvm::StringRef Path) {
185a9ac8606Spatrick   StringRef Name = llvm::sys::path::filename(Path);
186a9ac8606Spatrick   return Name.substr(0, Name.find('.', 1));
187a9ac8606Spatrick }
188a9ac8606Spatrick 
189e5dd7070Spatrick } // anonymous namespace
190e5dd7070Spatrick 
IncludeCategoryManager(const IncludeStyle & Style,StringRef FileName)191e5dd7070Spatrick IncludeCategoryManager::IncludeCategoryManager(const IncludeStyle &Style,
192e5dd7070Spatrick                                                StringRef FileName)
193e5dd7070Spatrick     : Style(Style), FileName(FileName) {
194a9ac8606Spatrick   for (const auto &Category : Style.IncludeCategories) {
195a9ac8606Spatrick     CategoryRegexs.emplace_back(Category.Regex, Category.RegexIsCaseSensitive
196a9ac8606Spatrick                                                     ? llvm::Regex::NoFlags
197a9ac8606Spatrick                                                     : llvm::Regex::IgnoreCase);
198a9ac8606Spatrick   }
199e5dd7070Spatrick   IsMainFile = FileName.endswith(".c") || FileName.endswith(".cc") ||
200e5dd7070Spatrick                FileName.endswith(".cpp") || FileName.endswith(".c++") ||
201e5dd7070Spatrick                FileName.endswith(".cxx") || FileName.endswith(".m") ||
202e5dd7070Spatrick                FileName.endswith(".mm");
203e5dd7070Spatrick   if (!Style.IncludeIsMainSourceRegex.empty()) {
204e5dd7070Spatrick     llvm::Regex MainFileRegex(Style.IncludeIsMainSourceRegex);
205e5dd7070Spatrick     IsMainFile |= MainFileRegex.match(FileName);
206e5dd7070Spatrick   }
207e5dd7070Spatrick }
208e5dd7070Spatrick 
getIncludePriority(StringRef IncludeName,bool CheckMainHeader) const209e5dd7070Spatrick int IncludeCategoryManager::getIncludePriority(StringRef IncludeName,
210e5dd7070Spatrick                                                bool CheckMainHeader) const {
211e5dd7070Spatrick   int Ret = INT_MAX;
212e5dd7070Spatrick   for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i)
213e5dd7070Spatrick     if (CategoryRegexs[i].match(IncludeName)) {
214e5dd7070Spatrick       Ret = Style.IncludeCategories[i].Priority;
215e5dd7070Spatrick       break;
216e5dd7070Spatrick     }
217e5dd7070Spatrick   if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
218e5dd7070Spatrick     Ret = 0;
219e5dd7070Spatrick   return Ret;
220e5dd7070Spatrick }
221e5dd7070Spatrick 
getSortIncludePriority(StringRef IncludeName,bool CheckMainHeader) const222e5dd7070Spatrick int IncludeCategoryManager::getSortIncludePriority(StringRef IncludeName,
223e5dd7070Spatrick                                                    bool CheckMainHeader) const {
224e5dd7070Spatrick   int Ret = INT_MAX;
225e5dd7070Spatrick   for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i)
226e5dd7070Spatrick     if (CategoryRegexs[i].match(IncludeName)) {
227e5dd7070Spatrick       Ret = Style.IncludeCategories[i].SortPriority;
228e5dd7070Spatrick       if (Ret == 0)
229e5dd7070Spatrick         Ret = Style.IncludeCategories[i].Priority;
230e5dd7070Spatrick       break;
231e5dd7070Spatrick     }
232e5dd7070Spatrick   if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
233e5dd7070Spatrick     Ret = 0;
234e5dd7070Spatrick   return Ret;
235e5dd7070Spatrick }
isMainHeader(StringRef IncludeName) const236e5dd7070Spatrick bool IncludeCategoryManager::isMainHeader(StringRef IncludeName) const {
237e5dd7070Spatrick   if (!IncludeName.startswith("\""))
238e5dd7070Spatrick     return false;
239a9ac8606Spatrick 
240a9ac8606Spatrick   IncludeName =
241a9ac8606Spatrick       IncludeName.drop_front(1).drop_back(1); // remove the surrounding "" or <>
242a9ac8606Spatrick   // Not matchingStem: implementation files may have compound extensions but
243a9ac8606Spatrick   // headers may not.
244a9ac8606Spatrick   StringRef HeaderStem = llvm::sys::path::stem(IncludeName);
245a9ac8606Spatrick   StringRef FileStem = llvm::sys::path::stem(FileName); // foo.cu for foo.cu.cc
246a9ac8606Spatrick   StringRef MatchingFileStem = matchingStem(FileName);  // foo for foo.cu.cc
247a9ac8606Spatrick   // main-header examples:
248a9ac8606Spatrick   //  1) foo.h => foo.cc
249a9ac8606Spatrick   //  2) foo.h => foo.cu.cc
250a9ac8606Spatrick   //  3) foo.proto.h => foo.proto.cc
251a9ac8606Spatrick   //
252a9ac8606Spatrick   // non-main-header examples:
253a9ac8606Spatrick   //  1) foo.h => bar.cc
254a9ac8606Spatrick   //  2) foo.proto.h => foo.cc
255a9ac8606Spatrick   StringRef Matching;
256a9ac8606Spatrick   if (MatchingFileStem.startswith_insensitive(HeaderStem))
257a9ac8606Spatrick     Matching = MatchingFileStem; // example 1), 2)
258a9ac8606Spatrick   else if (FileStem.equals_insensitive(HeaderStem))
259a9ac8606Spatrick     Matching = FileStem; // example 3)
260a9ac8606Spatrick   if (!Matching.empty()) {
261e5dd7070Spatrick     llvm::Regex MainIncludeRegex(HeaderStem.str() + Style.IncludeIsMainRegex,
262e5dd7070Spatrick                                  llvm::Regex::IgnoreCase);
263a9ac8606Spatrick     if (MainIncludeRegex.match(Matching))
264e5dd7070Spatrick       return true;
265e5dd7070Spatrick   }
266e5dd7070Spatrick   return false;
267e5dd7070Spatrick }
268e5dd7070Spatrick 
269*12c85518Srobert const llvm::Regex HeaderIncludes::IncludeRegex(IncludeRegexPattern);
270*12c85518Srobert 
HeaderIncludes(StringRef FileName,StringRef Code,const IncludeStyle & Style)271e5dd7070Spatrick HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
272e5dd7070Spatrick                                const IncludeStyle &Style)
273e5dd7070Spatrick     : FileName(FileName), Code(Code), FirstIncludeOffset(-1),
274e5dd7070Spatrick       MinInsertOffset(
275e5dd7070Spatrick           getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style)),
276e5dd7070Spatrick       MaxInsertOffset(MinInsertOffset +
277e5dd7070Spatrick                       getMaxHeaderInsertionOffset(
278e5dd7070Spatrick                           FileName, Code.drop_front(MinInsertOffset), Style)),
279*12c85518Srobert       Categories(Style, FileName) {
280e5dd7070Spatrick   // Add 0 for main header and INT_MAX for headers that are not in any
281e5dd7070Spatrick   // category.
282e5dd7070Spatrick   Priorities = {0, INT_MAX};
283e5dd7070Spatrick   for (const auto &Category : Style.IncludeCategories)
284e5dd7070Spatrick     Priorities.insert(Category.Priority);
285e5dd7070Spatrick   SmallVector<StringRef, 32> Lines;
286e5dd7070Spatrick   Code.drop_front(MinInsertOffset).split(Lines, "\n");
287e5dd7070Spatrick 
288e5dd7070Spatrick   unsigned Offset = MinInsertOffset;
289e5dd7070Spatrick   unsigned NextLineOffset;
290e5dd7070Spatrick   SmallVector<StringRef, 4> Matches;
291e5dd7070Spatrick   for (auto Line : Lines) {
292e5dd7070Spatrick     NextLineOffset = std::min(Code.size(), Offset + Line.size() + 1);
293e5dd7070Spatrick     if (IncludeRegex.match(Line, &Matches)) {
294e5dd7070Spatrick       // If this is the last line without trailing newline, we need to make
295e5dd7070Spatrick       // sure we don't delete across the file boundary.
296e5dd7070Spatrick       addExistingInclude(
297e5dd7070Spatrick           Include(Matches[2],
298e5dd7070Spatrick                   tooling::Range(
299*12c85518Srobert                       Offset, std::min(Line.size() + 1, Code.size() - Offset)),
300*12c85518Srobert                   Matches[1] == "import" ? tooling::IncludeDirective::Import
301*12c85518Srobert                                          : tooling::IncludeDirective::Include),
302e5dd7070Spatrick           NextLineOffset);
303e5dd7070Spatrick     }
304e5dd7070Spatrick     Offset = NextLineOffset;
305e5dd7070Spatrick   }
306e5dd7070Spatrick 
307e5dd7070Spatrick   // Populate CategoryEndOfssets:
308e5dd7070Spatrick   // - Ensure that CategoryEndOffset[Highest] is always populated.
309e5dd7070Spatrick   // - If CategoryEndOffset[Priority] isn't set, use the next higher value
310e5dd7070Spatrick   // that is set, up to CategoryEndOffset[Highest].
311e5dd7070Spatrick   auto Highest = Priorities.begin();
312e5dd7070Spatrick   if (CategoryEndOffsets.find(*Highest) == CategoryEndOffsets.end()) {
313e5dd7070Spatrick     if (FirstIncludeOffset >= 0)
314e5dd7070Spatrick       CategoryEndOffsets[*Highest] = FirstIncludeOffset;
315e5dd7070Spatrick     else
316e5dd7070Spatrick       CategoryEndOffsets[*Highest] = MinInsertOffset;
317e5dd7070Spatrick   }
318e5dd7070Spatrick   // By this point, CategoryEndOffset[Highest] is always set appropriately:
319e5dd7070Spatrick   //  - to an appropriate location before/after existing #includes, or
320e5dd7070Spatrick   //  - to right after the header guard, or
321e5dd7070Spatrick   //  - to the beginning of the file.
322e5dd7070Spatrick   for (auto I = ++Priorities.begin(), E = Priorities.end(); I != E; ++I)
323e5dd7070Spatrick     if (CategoryEndOffsets.find(*I) == CategoryEndOffsets.end())
324e5dd7070Spatrick       CategoryEndOffsets[*I] = CategoryEndOffsets[*std::prev(I)];
325e5dd7070Spatrick }
326e5dd7070Spatrick 
327e5dd7070Spatrick // \p Offset: the start of the line following this include directive.
addExistingInclude(Include IncludeToAdd,unsigned NextLineOffset)328e5dd7070Spatrick void HeaderIncludes::addExistingInclude(Include IncludeToAdd,
329e5dd7070Spatrick                                         unsigned NextLineOffset) {
330e5dd7070Spatrick   auto Iter =
331e5dd7070Spatrick       ExistingIncludes.try_emplace(trimInclude(IncludeToAdd.Name)).first;
332e5dd7070Spatrick   Iter->second.push_back(std::move(IncludeToAdd));
333e5dd7070Spatrick   auto &CurInclude = Iter->second.back();
334e5dd7070Spatrick   // The header name with quotes or angle brackets.
335e5dd7070Spatrick   // Only record the offset of current #include if we can insert after it.
336e5dd7070Spatrick   if (CurInclude.R.getOffset() <= MaxInsertOffset) {
337e5dd7070Spatrick     int Priority = Categories.getIncludePriority(
338e5dd7070Spatrick         CurInclude.Name, /*CheckMainHeader=*/FirstIncludeOffset < 0);
339e5dd7070Spatrick     CategoryEndOffsets[Priority] = NextLineOffset;
340e5dd7070Spatrick     IncludesByPriority[Priority].push_back(&CurInclude);
341e5dd7070Spatrick     if (FirstIncludeOffset < 0)
342e5dd7070Spatrick       FirstIncludeOffset = CurInclude.R.getOffset();
343e5dd7070Spatrick   }
344e5dd7070Spatrick }
345e5dd7070Spatrick 
346*12c85518Srobert std::optional<tooling::Replacement>
insert(llvm::StringRef IncludeName,bool IsAngled,IncludeDirective Directive) const347*12c85518Srobert HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled,
348*12c85518Srobert                        IncludeDirective Directive) const {
349e5dd7070Spatrick   assert(IncludeName == trimInclude(IncludeName));
350e5dd7070Spatrick   // If a <header> ("header") already exists in code, "header" (<header>) with
351*12c85518Srobert   // different quotation and/or directive will still be inserted.
352e5dd7070Spatrick   // FIXME: figure out if this is the best behavior.
353e5dd7070Spatrick   auto It = ExistingIncludes.find(IncludeName);
354*12c85518Srobert   if (It != ExistingIncludes.end()) {
355e5dd7070Spatrick     for (const auto &Inc : It->second)
356*12c85518Srobert       if (Inc.Directive == Directive &&
357*12c85518Srobert           ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
358*12c85518Srobert            (!IsAngled && StringRef(Inc.Name).startswith("\""))))
359*12c85518Srobert         return std::nullopt;
360*12c85518Srobert   }
361e5dd7070Spatrick   std::string Quoted =
362ec727ea7Spatrick       std::string(llvm::formatv(IsAngled ? "<{0}>" : "\"{0}\"", IncludeName));
363e5dd7070Spatrick   StringRef QuotedName = Quoted;
364e5dd7070Spatrick   int Priority = Categories.getIncludePriority(
365e5dd7070Spatrick       QuotedName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
366e5dd7070Spatrick   auto CatOffset = CategoryEndOffsets.find(Priority);
367e5dd7070Spatrick   assert(CatOffset != CategoryEndOffsets.end());
368e5dd7070Spatrick   unsigned InsertOffset = CatOffset->second; // Fall back offset
369e5dd7070Spatrick   auto Iter = IncludesByPriority.find(Priority);
370e5dd7070Spatrick   if (Iter != IncludesByPriority.end()) {
371e5dd7070Spatrick     for (const auto *Inc : Iter->second) {
372e5dd7070Spatrick       if (QuotedName < Inc->Name) {
373e5dd7070Spatrick         InsertOffset = Inc->R.getOffset();
374e5dd7070Spatrick         break;
375e5dd7070Spatrick       }
376e5dd7070Spatrick     }
377e5dd7070Spatrick   }
378e5dd7070Spatrick   assert(InsertOffset <= Code.size());
379*12c85518Srobert   llvm::StringRef DirectiveSpelling =
380*12c85518Srobert       Directive == IncludeDirective::Include ? "include" : "import";
381ec727ea7Spatrick   std::string NewInclude =
382*12c85518Srobert       llvm::formatv("#{0} {1}\n", DirectiveSpelling, QuotedName);
383e5dd7070Spatrick   // When inserting headers at end of the code, also append '\n' to the code
384e5dd7070Spatrick   // if it does not end with '\n'.
385e5dd7070Spatrick   // FIXME: when inserting multiple #includes at the end of code, only one
386e5dd7070Spatrick   // newline should be added.
387e5dd7070Spatrick   if (InsertOffset == Code.size() && (!Code.empty() && Code.back() != '\n'))
388e5dd7070Spatrick     NewInclude = "\n" + NewInclude;
389e5dd7070Spatrick   return tooling::Replacement(FileName, InsertOffset, 0, NewInclude);
390e5dd7070Spatrick }
391e5dd7070Spatrick 
remove(llvm::StringRef IncludeName,bool IsAngled) const392e5dd7070Spatrick tooling::Replacements HeaderIncludes::remove(llvm::StringRef IncludeName,
393e5dd7070Spatrick                                              bool IsAngled) const {
394e5dd7070Spatrick   assert(IncludeName == trimInclude(IncludeName));
395e5dd7070Spatrick   tooling::Replacements Result;
396e5dd7070Spatrick   auto Iter = ExistingIncludes.find(IncludeName);
397e5dd7070Spatrick   if (Iter == ExistingIncludes.end())
398e5dd7070Spatrick     return Result;
399e5dd7070Spatrick   for (const auto &Inc : Iter->second) {
400e5dd7070Spatrick     if ((IsAngled && StringRef(Inc.Name).startswith("\"")) ||
401e5dd7070Spatrick         (!IsAngled && StringRef(Inc.Name).startswith("<")))
402e5dd7070Spatrick       continue;
403e5dd7070Spatrick     llvm::Error Err = Result.add(tooling::Replacement(
404e5dd7070Spatrick         FileName, Inc.R.getOffset(), Inc.R.getLength(), ""));
405e5dd7070Spatrick     if (Err) {
406e5dd7070Spatrick       auto ErrMsg = "Unexpected conflicts in #include deletions: " +
407e5dd7070Spatrick                     llvm::toString(std::move(Err));
408e5dd7070Spatrick       llvm_unreachable(ErrMsg.c_str());
409e5dd7070Spatrick     }
410e5dd7070Spatrick   }
411e5dd7070Spatrick   return Result;
412e5dd7070Spatrick }
413e5dd7070Spatrick 
414e5dd7070Spatrick } // namespace tooling
415e5dd7070Spatrick } // namespace clang
416