xref: /netbsd-src/external/apache2/llvm/dist/clang/lib/Tooling/Inclusions/HeaderIncludes.cpp (revision e038c9c4676b0f19b1b7dd08a940c6ed64a6d5ae)
17330f729Sjoerg //===--- HeaderIncludes.cpp - Insert/Delete #includes --*- C++ -*----------===//
27330f729Sjoerg //
37330f729Sjoerg // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
47330f729Sjoerg // See https://llvm.org/LICENSE.txt for license information.
57330f729Sjoerg // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
67330f729Sjoerg //
77330f729Sjoerg //===----------------------------------------------------------------------===//
87330f729Sjoerg 
97330f729Sjoerg #include "clang/Tooling/Inclusions/HeaderIncludes.h"
10*e038c9c4Sjoerg #include "clang/Basic/FileManager.h"
117330f729Sjoerg #include "clang/Basic/SourceManager.h"
127330f729Sjoerg #include "clang/Lex/Lexer.h"
137330f729Sjoerg #include "llvm/ADT/Optional.h"
147330f729Sjoerg #include "llvm/Support/FormatVariadic.h"
15*e038c9c4Sjoerg #include "llvm/Support/Path.h"
167330f729Sjoerg 
177330f729Sjoerg namespace clang {
187330f729Sjoerg namespace tooling {
197330f729Sjoerg namespace {
207330f729Sjoerg 
createLangOpts()217330f729Sjoerg LangOptions createLangOpts() {
227330f729Sjoerg   LangOptions LangOpts;
237330f729Sjoerg   LangOpts.CPlusPlus = 1;
247330f729Sjoerg   LangOpts.CPlusPlus11 = 1;
257330f729Sjoerg   LangOpts.CPlusPlus14 = 1;
267330f729Sjoerg   LangOpts.LineComment = 1;
277330f729Sjoerg   LangOpts.CXXOperatorNames = 1;
287330f729Sjoerg   LangOpts.Bool = 1;
297330f729Sjoerg   LangOpts.ObjC = 1;
307330f729Sjoerg   LangOpts.MicrosoftExt = 1;    // To get kw___try, kw___finally.
317330f729Sjoerg   LangOpts.DeclSpecKeyword = 1; // To get __declspec.
327330f729Sjoerg   LangOpts.WChar = 1;           // To get wchar_t
337330f729Sjoerg   return LangOpts;
347330f729Sjoerg }
357330f729Sjoerg 
367330f729Sjoerg // Returns the offset after skipping a sequence of tokens, matched by \p
377330f729Sjoerg // GetOffsetAfterSequence, from the start of the code.
387330f729Sjoerg // \p GetOffsetAfterSequence should be a function that matches a sequence of
397330f729Sjoerg // 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)407330f729Sjoerg unsigned getOffsetAfterTokenSequence(
417330f729Sjoerg     StringRef FileName, StringRef Code, const IncludeStyle &Style,
427330f729Sjoerg     llvm::function_ref<unsigned(const SourceManager &, Lexer &, Token &)>
437330f729Sjoerg         GetOffsetAfterSequence) {
447330f729Sjoerg   SourceManagerForFile VirtualSM(FileName, Code);
457330f729Sjoerg   SourceManager &SM = VirtualSM.get();
46*e038c9c4Sjoerg   Lexer Lex(SM.getMainFileID(), SM.getBufferOrFake(SM.getMainFileID()), SM,
477330f729Sjoerg             createLangOpts());
487330f729Sjoerg   Token Tok;
497330f729Sjoerg   // Get the first token.
507330f729Sjoerg   Lex.LexFromRawLexer(Tok);
517330f729Sjoerg   return GetOffsetAfterSequence(SM, Lex, Tok);
527330f729Sjoerg }
537330f729Sjoerg 
547330f729Sjoerg // Check if a sequence of tokens is like "#<Name> <raw_identifier>". If it is,
557330f729Sjoerg // \p Tok will be the token after this directive; otherwise, it can be any token
567330f729Sjoerg // after the given \p Tok (including \p Tok). If \p RawIDName is provided, the
577330f729Sjoerg // (second) raw_identifier name is checked.
checkAndConsumeDirectiveWithName(Lexer & Lex,StringRef Name,Token & Tok,llvm::Optional<StringRef> RawIDName=llvm::None)587330f729Sjoerg bool checkAndConsumeDirectiveWithName(
597330f729Sjoerg     Lexer &Lex, StringRef Name, Token &Tok,
607330f729Sjoerg     llvm::Optional<StringRef> RawIDName = llvm::None) {
617330f729Sjoerg   bool Matched = Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
627330f729Sjoerg                  Tok.is(tok::raw_identifier) &&
637330f729Sjoerg                  Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) &&
647330f729Sjoerg                  Tok.is(tok::raw_identifier) &&
657330f729Sjoerg                  (!RawIDName || Tok.getRawIdentifier() == *RawIDName);
667330f729Sjoerg   if (Matched)
677330f729Sjoerg     Lex.LexFromRawLexer(Tok);
687330f729Sjoerg   return Matched;
697330f729Sjoerg }
707330f729Sjoerg 
skipComments(Lexer & Lex,Token & Tok)717330f729Sjoerg void skipComments(Lexer &Lex, Token &Tok) {
727330f729Sjoerg   while (Tok.is(tok::comment))
737330f729Sjoerg     if (Lex.LexFromRawLexer(Tok))
747330f729Sjoerg       return;
757330f729Sjoerg }
767330f729Sjoerg 
777330f729Sjoerg // Returns the offset after header guard directives and any comments
787330f729Sjoerg // before/after header guards (e.g. #ifndef/#define pair, #pragma once). If no
797330f729Sjoerg // header guard is present in the code, this will return the offset after
807330f729Sjoerg // skipping all comments from the start of the code.
getOffsetAfterHeaderGuardsAndComments(StringRef FileName,StringRef Code,const IncludeStyle & Style)817330f729Sjoerg unsigned getOffsetAfterHeaderGuardsAndComments(StringRef FileName,
827330f729Sjoerg                                                StringRef Code,
837330f729Sjoerg                                                const IncludeStyle &Style) {
847330f729Sjoerg   // \p Consume returns location after header guard or 0 if no header guard is
857330f729Sjoerg   // found.
867330f729Sjoerg   auto ConsumeHeaderGuardAndComment =
877330f729Sjoerg       [&](std::function<unsigned(const SourceManager &SM, Lexer &Lex,
887330f729Sjoerg                                  Token Tok)>
897330f729Sjoerg               Consume) {
907330f729Sjoerg         return getOffsetAfterTokenSequence(
917330f729Sjoerg             FileName, Code, Style,
927330f729Sjoerg             [&Consume](const SourceManager &SM, Lexer &Lex, Token Tok) {
937330f729Sjoerg               skipComments(Lex, Tok);
947330f729Sjoerg               unsigned InitialOffset = SM.getFileOffset(Tok.getLocation());
957330f729Sjoerg               return std::max(InitialOffset, Consume(SM, Lex, Tok));
967330f729Sjoerg             });
977330f729Sjoerg       };
987330f729Sjoerg   return std::max(
997330f729Sjoerg       // #ifndef/#define
1007330f729Sjoerg       ConsumeHeaderGuardAndComment(
1017330f729Sjoerg           [](const SourceManager &SM, Lexer &Lex, Token Tok) -> unsigned {
1027330f729Sjoerg             if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) {
1037330f729Sjoerg               skipComments(Lex, Tok);
104*e038c9c4Sjoerg               if (checkAndConsumeDirectiveWithName(Lex, "define", Tok) &&
105*e038c9c4Sjoerg                   Tok.isAtStartOfLine())
1067330f729Sjoerg                 return SM.getFileOffset(Tok.getLocation());
1077330f729Sjoerg             }
1087330f729Sjoerg             return 0;
1097330f729Sjoerg           }),
1107330f729Sjoerg       // #pragma once
1117330f729Sjoerg       ConsumeHeaderGuardAndComment(
1127330f729Sjoerg           [](const SourceManager &SM, Lexer &Lex, Token Tok) -> unsigned {
1137330f729Sjoerg             if (checkAndConsumeDirectiveWithName(Lex, "pragma", Tok,
1147330f729Sjoerg                                                  StringRef("once")))
1157330f729Sjoerg               return SM.getFileOffset(Tok.getLocation());
1167330f729Sjoerg             return 0;
1177330f729Sjoerg           }));
1187330f729Sjoerg }
1197330f729Sjoerg 
1207330f729Sjoerg // Check if a sequence of tokens is like
1217330f729Sjoerg //    "#include ("header.h" | <header.h>)".
1227330f729Sjoerg // If it is, \p Tok will be the token after this directive; otherwise, it can be
1237330f729Sjoerg // any token after the given \p Tok (including \p Tok).
checkAndConsumeInclusiveDirective(Lexer & Lex,Token & Tok)1247330f729Sjoerg bool checkAndConsumeInclusiveDirective(Lexer &Lex, Token &Tok) {
1257330f729Sjoerg   auto Matched = [&]() {
1267330f729Sjoerg     Lex.LexFromRawLexer(Tok);
1277330f729Sjoerg     return true;
1287330f729Sjoerg   };
1297330f729Sjoerg   if (Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
1307330f729Sjoerg       Tok.is(tok::raw_identifier) && Tok.getRawIdentifier() == "include") {
1317330f729Sjoerg     if (Lex.LexFromRawLexer(Tok))
1327330f729Sjoerg       return false;
1337330f729Sjoerg     if (Tok.is(tok::string_literal))
1347330f729Sjoerg       return Matched();
1357330f729Sjoerg     if (Tok.is(tok::less)) {
1367330f729Sjoerg       while (!Lex.LexFromRawLexer(Tok) && Tok.isNot(tok::greater)) {
1377330f729Sjoerg       }
1387330f729Sjoerg       if (Tok.is(tok::greater))
1397330f729Sjoerg         return Matched();
1407330f729Sjoerg     }
1417330f729Sjoerg   }
1427330f729Sjoerg   return false;
1437330f729Sjoerg }
1447330f729Sjoerg 
1457330f729Sjoerg // Returns the offset of the last #include directive after which a new
1467330f729Sjoerg // #include can be inserted. This ignores #include's after the #include block(s)
1477330f729Sjoerg // in the beginning of a file to avoid inserting headers into code sections
1487330f729Sjoerg // where new #include's should not be added by default.
1497330f729Sjoerg // These code sections include:
1507330f729Sjoerg //      - raw string literals (containing #include).
1517330f729Sjoerg //      - #if blocks.
1527330f729Sjoerg //      - Special #include's among declarations (e.g. functions).
1537330f729Sjoerg //
1547330f729Sjoerg // If no #include after which a new #include can be inserted, this returns the
1557330f729Sjoerg // offset after skipping all comments from the start of the code.
1567330f729Sjoerg // Inserting after an #include is not allowed if it comes after code that is not
1577330f729Sjoerg // #include (e.g. pre-processing directive that is not #include, declarations).
getMaxHeaderInsertionOffset(StringRef FileName,StringRef Code,const IncludeStyle & Style)1587330f729Sjoerg unsigned getMaxHeaderInsertionOffset(StringRef FileName, StringRef Code,
1597330f729Sjoerg                                      const IncludeStyle &Style) {
1607330f729Sjoerg   return getOffsetAfterTokenSequence(
1617330f729Sjoerg       FileName, Code, Style,
1627330f729Sjoerg       [](const SourceManager &SM, Lexer &Lex, Token Tok) {
1637330f729Sjoerg         skipComments(Lex, Tok);
1647330f729Sjoerg         unsigned MaxOffset = SM.getFileOffset(Tok.getLocation());
1657330f729Sjoerg         while (checkAndConsumeInclusiveDirective(Lex, Tok))
1667330f729Sjoerg           MaxOffset = SM.getFileOffset(Tok.getLocation());
1677330f729Sjoerg         return MaxOffset;
1687330f729Sjoerg       });
1697330f729Sjoerg }
1707330f729Sjoerg 
trimInclude(StringRef IncludeName)1717330f729Sjoerg inline StringRef trimInclude(StringRef IncludeName) {
1727330f729Sjoerg   return IncludeName.trim("\"<>");
1737330f729Sjoerg }
1747330f729Sjoerg 
1757330f729Sjoerg const char IncludeRegexPattern[] =
1767330f729Sjoerg     R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))";
1777330f729Sjoerg 
178*e038c9c4Sjoerg // The filename of Path excluding extension.
179*e038c9c4Sjoerg // Used to match implementation with headers, this differs from sys::path::stem:
180*e038c9c4Sjoerg //  - in names with multiple dots (foo.cu.cc) it terminates at the *first*
181*e038c9c4Sjoerg //  - an empty stem is never returned: /foo/.bar.x => .bar
182*e038c9c4Sjoerg //  - we don't bother to handle . and .. specially
matchingStem(llvm::StringRef Path)183*e038c9c4Sjoerg StringRef matchingStem(llvm::StringRef Path) {
184*e038c9c4Sjoerg   StringRef Name = llvm::sys::path::filename(Path);
185*e038c9c4Sjoerg   return Name.substr(0, Name.find('.', 1));
186*e038c9c4Sjoerg }
187*e038c9c4Sjoerg 
1887330f729Sjoerg } // anonymous namespace
1897330f729Sjoerg 
IncludeCategoryManager(const IncludeStyle & Style,StringRef FileName)1907330f729Sjoerg IncludeCategoryManager::IncludeCategoryManager(const IncludeStyle &Style,
1917330f729Sjoerg                                                StringRef FileName)
1927330f729Sjoerg     : Style(Style), FileName(FileName) {
193*e038c9c4Sjoerg   for (const auto &Category : Style.IncludeCategories) {
194*e038c9c4Sjoerg     CategoryRegexs.emplace_back(Category.Regex, Category.RegexIsCaseSensitive
195*e038c9c4Sjoerg                                                     ? llvm::Regex::NoFlags
196*e038c9c4Sjoerg                                                     : llvm::Regex::IgnoreCase);
197*e038c9c4Sjoerg   }
1987330f729Sjoerg   IsMainFile = FileName.endswith(".c") || FileName.endswith(".cc") ||
1997330f729Sjoerg                FileName.endswith(".cpp") || FileName.endswith(".c++") ||
2007330f729Sjoerg                FileName.endswith(".cxx") || FileName.endswith(".m") ||
2017330f729Sjoerg                FileName.endswith(".mm");
202*e038c9c4Sjoerg   if (!Style.IncludeIsMainSourceRegex.empty()) {
203*e038c9c4Sjoerg     llvm::Regex MainFileRegex(Style.IncludeIsMainSourceRegex);
204*e038c9c4Sjoerg     IsMainFile |= MainFileRegex.match(FileName);
205*e038c9c4Sjoerg   }
2067330f729Sjoerg }
2077330f729Sjoerg 
getIncludePriority(StringRef IncludeName,bool CheckMainHeader) const2087330f729Sjoerg int IncludeCategoryManager::getIncludePriority(StringRef IncludeName,
2097330f729Sjoerg                                                bool CheckMainHeader) const {
2107330f729Sjoerg   int Ret = INT_MAX;
2117330f729Sjoerg   for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i)
2127330f729Sjoerg     if (CategoryRegexs[i].match(IncludeName)) {
2137330f729Sjoerg       Ret = Style.IncludeCategories[i].Priority;
2147330f729Sjoerg       break;
2157330f729Sjoerg     }
2167330f729Sjoerg   if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
2177330f729Sjoerg     Ret = 0;
2187330f729Sjoerg   return Ret;
2197330f729Sjoerg }
2207330f729Sjoerg 
getSortIncludePriority(StringRef IncludeName,bool CheckMainHeader) const2217330f729Sjoerg int IncludeCategoryManager::getSortIncludePriority(StringRef IncludeName,
2227330f729Sjoerg                                                    bool CheckMainHeader) const {
2237330f729Sjoerg   int Ret = INT_MAX;
2247330f729Sjoerg   for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i)
2257330f729Sjoerg     if (CategoryRegexs[i].match(IncludeName)) {
2267330f729Sjoerg       Ret = Style.IncludeCategories[i].SortPriority;
2277330f729Sjoerg       if (Ret == 0)
2287330f729Sjoerg         Ret = Style.IncludeCategories[i].Priority;
2297330f729Sjoerg       break;
2307330f729Sjoerg     }
2317330f729Sjoerg   if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
2327330f729Sjoerg     Ret = 0;
2337330f729Sjoerg   return Ret;
2347330f729Sjoerg }
isMainHeader(StringRef IncludeName) const2357330f729Sjoerg bool IncludeCategoryManager::isMainHeader(StringRef IncludeName) const {
2367330f729Sjoerg   if (!IncludeName.startswith("\""))
2377330f729Sjoerg     return false;
238*e038c9c4Sjoerg 
239*e038c9c4Sjoerg   IncludeName =
240*e038c9c4Sjoerg       IncludeName.drop_front(1).drop_back(1); // remove the surrounding "" or <>
241*e038c9c4Sjoerg   // Not matchingStem: implementation files may have compound extensions but
242*e038c9c4Sjoerg   // headers may not.
243*e038c9c4Sjoerg   StringRef HeaderStem = llvm::sys::path::stem(IncludeName);
244*e038c9c4Sjoerg   StringRef FileStem = llvm::sys::path::stem(FileName); // foo.cu for foo.cu.cc
245*e038c9c4Sjoerg   StringRef MatchingFileStem = matchingStem(FileName);  // foo for foo.cu.cc
246*e038c9c4Sjoerg   // main-header examples:
247*e038c9c4Sjoerg   //  1) foo.h => foo.cc
248*e038c9c4Sjoerg   //  2) foo.h => foo.cu.cc
249*e038c9c4Sjoerg   //  3) foo.proto.h => foo.proto.cc
250*e038c9c4Sjoerg   //
251*e038c9c4Sjoerg   // non-main-header examples:
252*e038c9c4Sjoerg   //  1) foo.h => bar.cc
253*e038c9c4Sjoerg   //  2) foo.proto.h => foo.cc
254*e038c9c4Sjoerg   StringRef Matching;
255*e038c9c4Sjoerg   if (MatchingFileStem.startswith_lower(HeaderStem))
256*e038c9c4Sjoerg     Matching = MatchingFileStem; // example 1), 2)
257*e038c9c4Sjoerg   else if (FileStem.equals_lower(HeaderStem))
258*e038c9c4Sjoerg     Matching = FileStem; // example 3)
259*e038c9c4Sjoerg   if (!Matching.empty()) {
2607330f729Sjoerg     llvm::Regex MainIncludeRegex(HeaderStem.str() + Style.IncludeIsMainRegex,
2617330f729Sjoerg                                  llvm::Regex::IgnoreCase);
262*e038c9c4Sjoerg     if (MainIncludeRegex.match(Matching))
2637330f729Sjoerg       return true;
2647330f729Sjoerg   }
2657330f729Sjoerg   return false;
2667330f729Sjoerg }
2677330f729Sjoerg 
HeaderIncludes(StringRef FileName,StringRef Code,const IncludeStyle & Style)2687330f729Sjoerg HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
2697330f729Sjoerg                                const IncludeStyle &Style)
2707330f729Sjoerg     : FileName(FileName), Code(Code), FirstIncludeOffset(-1),
2717330f729Sjoerg       MinInsertOffset(
2727330f729Sjoerg           getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style)),
2737330f729Sjoerg       MaxInsertOffset(MinInsertOffset +
2747330f729Sjoerg                       getMaxHeaderInsertionOffset(
2757330f729Sjoerg                           FileName, Code.drop_front(MinInsertOffset), Style)),
2767330f729Sjoerg       Categories(Style, FileName),
2777330f729Sjoerg       IncludeRegex(llvm::Regex(IncludeRegexPattern)) {
2787330f729Sjoerg   // Add 0 for main header and INT_MAX for headers that are not in any
2797330f729Sjoerg   // category.
2807330f729Sjoerg   Priorities = {0, INT_MAX};
2817330f729Sjoerg   for (const auto &Category : Style.IncludeCategories)
2827330f729Sjoerg     Priorities.insert(Category.Priority);
2837330f729Sjoerg   SmallVector<StringRef, 32> Lines;
2847330f729Sjoerg   Code.drop_front(MinInsertOffset).split(Lines, "\n");
2857330f729Sjoerg 
2867330f729Sjoerg   unsigned Offset = MinInsertOffset;
2877330f729Sjoerg   unsigned NextLineOffset;
2887330f729Sjoerg   SmallVector<StringRef, 4> Matches;
2897330f729Sjoerg   for (auto Line : Lines) {
2907330f729Sjoerg     NextLineOffset = std::min(Code.size(), Offset + Line.size() + 1);
2917330f729Sjoerg     if (IncludeRegex.match(Line, &Matches)) {
2927330f729Sjoerg       // If this is the last line without trailing newline, we need to make
2937330f729Sjoerg       // sure we don't delete across the file boundary.
2947330f729Sjoerg       addExistingInclude(
2957330f729Sjoerg           Include(Matches[2],
2967330f729Sjoerg                   tooling::Range(
2977330f729Sjoerg                       Offset, std::min(Line.size() + 1, Code.size() - Offset))),
2987330f729Sjoerg           NextLineOffset);
2997330f729Sjoerg     }
3007330f729Sjoerg     Offset = NextLineOffset;
3017330f729Sjoerg   }
3027330f729Sjoerg 
3037330f729Sjoerg   // Populate CategoryEndOfssets:
3047330f729Sjoerg   // - Ensure that CategoryEndOffset[Highest] is always populated.
3057330f729Sjoerg   // - If CategoryEndOffset[Priority] isn't set, use the next higher value
3067330f729Sjoerg   // that is set, up to CategoryEndOffset[Highest].
3077330f729Sjoerg   auto Highest = Priorities.begin();
3087330f729Sjoerg   if (CategoryEndOffsets.find(*Highest) == CategoryEndOffsets.end()) {
3097330f729Sjoerg     if (FirstIncludeOffset >= 0)
3107330f729Sjoerg       CategoryEndOffsets[*Highest] = FirstIncludeOffset;
3117330f729Sjoerg     else
3127330f729Sjoerg       CategoryEndOffsets[*Highest] = MinInsertOffset;
3137330f729Sjoerg   }
3147330f729Sjoerg   // By this point, CategoryEndOffset[Highest] is always set appropriately:
3157330f729Sjoerg   //  - to an appropriate location before/after existing #includes, or
3167330f729Sjoerg   //  - to right after the header guard, or
3177330f729Sjoerg   //  - to the beginning of the file.
3187330f729Sjoerg   for (auto I = ++Priorities.begin(), E = Priorities.end(); I != E; ++I)
3197330f729Sjoerg     if (CategoryEndOffsets.find(*I) == CategoryEndOffsets.end())
3207330f729Sjoerg       CategoryEndOffsets[*I] = CategoryEndOffsets[*std::prev(I)];
3217330f729Sjoerg }
3227330f729Sjoerg 
3237330f729Sjoerg // \p Offset: the start of the line following this include directive.
addExistingInclude(Include IncludeToAdd,unsigned NextLineOffset)3247330f729Sjoerg void HeaderIncludes::addExistingInclude(Include IncludeToAdd,
3257330f729Sjoerg                                         unsigned NextLineOffset) {
3267330f729Sjoerg   auto Iter =
3277330f729Sjoerg       ExistingIncludes.try_emplace(trimInclude(IncludeToAdd.Name)).first;
3287330f729Sjoerg   Iter->second.push_back(std::move(IncludeToAdd));
3297330f729Sjoerg   auto &CurInclude = Iter->second.back();
3307330f729Sjoerg   // The header name with quotes or angle brackets.
3317330f729Sjoerg   // Only record the offset of current #include if we can insert after it.
3327330f729Sjoerg   if (CurInclude.R.getOffset() <= MaxInsertOffset) {
3337330f729Sjoerg     int Priority = Categories.getIncludePriority(
3347330f729Sjoerg         CurInclude.Name, /*CheckMainHeader=*/FirstIncludeOffset < 0);
3357330f729Sjoerg     CategoryEndOffsets[Priority] = NextLineOffset;
3367330f729Sjoerg     IncludesByPriority[Priority].push_back(&CurInclude);
3377330f729Sjoerg     if (FirstIncludeOffset < 0)
3387330f729Sjoerg       FirstIncludeOffset = CurInclude.R.getOffset();
3397330f729Sjoerg   }
3407330f729Sjoerg }
3417330f729Sjoerg 
3427330f729Sjoerg llvm::Optional<tooling::Replacement>
insert(llvm::StringRef IncludeName,bool IsAngled) const3437330f729Sjoerg HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled) const {
3447330f729Sjoerg   assert(IncludeName == trimInclude(IncludeName));
3457330f729Sjoerg   // If a <header> ("header") already exists in code, "header" (<header>) with
3467330f729Sjoerg   // different quotation will still be inserted.
3477330f729Sjoerg   // FIXME: figure out if this is the best behavior.
3487330f729Sjoerg   auto It = ExistingIncludes.find(IncludeName);
3497330f729Sjoerg   if (It != ExistingIncludes.end())
3507330f729Sjoerg     for (const auto &Inc : It->second)
3517330f729Sjoerg       if ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
3527330f729Sjoerg           (!IsAngled && StringRef(Inc.Name).startswith("\"")))
3537330f729Sjoerg         return llvm::None;
3547330f729Sjoerg   std::string Quoted =
355*e038c9c4Sjoerg       std::string(llvm::formatv(IsAngled ? "<{0}>" : "\"{0}\"", IncludeName));
3567330f729Sjoerg   StringRef QuotedName = Quoted;
3577330f729Sjoerg   int Priority = Categories.getIncludePriority(
3587330f729Sjoerg       QuotedName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
3597330f729Sjoerg   auto CatOffset = CategoryEndOffsets.find(Priority);
3607330f729Sjoerg   assert(CatOffset != CategoryEndOffsets.end());
3617330f729Sjoerg   unsigned InsertOffset = CatOffset->second; // Fall back offset
3627330f729Sjoerg   auto Iter = IncludesByPriority.find(Priority);
3637330f729Sjoerg   if (Iter != IncludesByPriority.end()) {
3647330f729Sjoerg     for (const auto *Inc : Iter->second) {
3657330f729Sjoerg       if (QuotedName < Inc->Name) {
3667330f729Sjoerg         InsertOffset = Inc->R.getOffset();
3677330f729Sjoerg         break;
3687330f729Sjoerg       }
3697330f729Sjoerg     }
3707330f729Sjoerg   }
3717330f729Sjoerg   assert(InsertOffset <= Code.size());
372*e038c9c4Sjoerg   std::string NewInclude =
373*e038c9c4Sjoerg       std::string(llvm::formatv("#include {0}\n", QuotedName));
3747330f729Sjoerg   // When inserting headers at end of the code, also append '\n' to the code
3757330f729Sjoerg   // if it does not end with '\n'.
3767330f729Sjoerg   // FIXME: when inserting multiple #includes at the end of code, only one
3777330f729Sjoerg   // newline should be added.
3787330f729Sjoerg   if (InsertOffset == Code.size() && (!Code.empty() && Code.back() != '\n'))
3797330f729Sjoerg     NewInclude = "\n" + NewInclude;
3807330f729Sjoerg   return tooling::Replacement(FileName, InsertOffset, 0, NewInclude);
3817330f729Sjoerg }
3827330f729Sjoerg 
remove(llvm::StringRef IncludeName,bool IsAngled) const3837330f729Sjoerg tooling::Replacements HeaderIncludes::remove(llvm::StringRef IncludeName,
3847330f729Sjoerg                                              bool IsAngled) const {
3857330f729Sjoerg   assert(IncludeName == trimInclude(IncludeName));
3867330f729Sjoerg   tooling::Replacements Result;
3877330f729Sjoerg   auto Iter = ExistingIncludes.find(IncludeName);
3887330f729Sjoerg   if (Iter == ExistingIncludes.end())
3897330f729Sjoerg     return Result;
3907330f729Sjoerg   for (const auto &Inc : Iter->second) {
3917330f729Sjoerg     if ((IsAngled && StringRef(Inc.Name).startswith("\"")) ||
3927330f729Sjoerg         (!IsAngled && StringRef(Inc.Name).startswith("<")))
3937330f729Sjoerg       continue;
3947330f729Sjoerg     llvm::Error Err = Result.add(tooling::Replacement(
3957330f729Sjoerg         FileName, Inc.R.getOffset(), Inc.R.getLength(), ""));
3967330f729Sjoerg     if (Err) {
3977330f729Sjoerg       auto ErrMsg = "Unexpected conflicts in #include deletions: " +
3987330f729Sjoerg                     llvm::toString(std::move(Err));
3997330f729Sjoerg       llvm_unreachable(ErrMsg.c_str());
4007330f729Sjoerg     }
4017330f729Sjoerg   }
4027330f729Sjoerg   return Result;
4037330f729Sjoerg }
4047330f729Sjoerg 
4057330f729Sjoerg } // namespace tooling
4067330f729Sjoerg } // namespace clang
407