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