clang/lib/Tooling/Syntax/Tokens.cpp

   1 //===- Tokens.cpp - collect tokens from preprocessing ---------------------===//
   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 #include "clang/Tooling/Syntax/Tokens.h"
   9
  10 #include "clang/Basic/Diagnostic.h"
  11 #include "clang/Basic/IdentifierTable.h"
  12 #include "clang/Basic/LLVM.h"
  13 #include "clang/Basic/LangOptions.h"
  14 #include "clang/Basic/SourceLocation.h"
  15 #include "clang/Basic/SourceManager.h"
  16 #include "clang/Basic/TokenKinds.h"
  17 #include "clang/Lex/PPCallbacks.h"
  18 #include "clang/Lex/Preprocessor.h"
  19 #include "clang/Lex/Token.h"
  20 #include "llvm/ADT/ArrayRef.h"
  21 #include "llvm/ADT/STLExtras.h"
  22 #include "llvm/Support/Debug.h"
  23 #include "llvm/Support/ErrorHandling.h"
  24 #include "llvm/Support/FormatVariadic.h"
  25 #include "llvm/Support/raw_ostream.h"
  26 #include <algorithm>
  27 #include <cassert>
  28 #include <iterator>
  29 #include <optional>
  30 #include <string>
  31 #include <utility>
  32 #include <vector>
  33
  34 using namespace clang;
  35 using namespace clang::syntax;
  36
  37 namespace {
  38 // Finds the smallest consecutive subsuquence of Toks that covers R.
  39 llvm::ArrayRef<syntax::Token>
  40 getTokensCovering(llvm::ArrayRef<syntax::Token> Toks, SourceRange R,
  41                   const SourceManager &SM) {
  42   if (R.isInvalid())
  43     return {};
  44   const syntax::Token *Begin =
  45       llvm::partition_point(Toks, [&](const syntax::Token &T) {
  46         return SM.isBeforeInTranslationUnit(T.location(), R.getBegin());
  47       });
  48   const syntax::Token *End =
  49       llvm::partition_point(Toks, [&](const syntax::Token &T) {
  50         return !SM.isBeforeInTranslationUnit(R.getEnd(), T.location());
  51       });
  52   if (Begin > End)
  53     return {};
  54   return {Begin, End};
  55 }
  56
  57 // Finds the range within FID corresponding to expanded tokens [First, Last].
  58 // Prev precedes First and Next follows Last, these must *not* be included.
  59 // If no range satisfies the criteria, returns an invalid range.
  60 //
  61 // #define ID(x) x
  62 // ID(ID(ID(a1) a2))
  63 //          ~~       -> a1
  64 //              ~~   -> a2
  65 //       ~~~~~~~~~   -> a1 a2
  66 SourceRange spelledForExpandedSlow(SourceLocation First, SourceLocation Last,
  67                                    SourceLocation Prev, SourceLocation Next,
  68                                    FileID TargetFile,
  69                                    const SourceManager &SM) {
  70   // There are two main parts to this algorithm:
  71   //  - identifying which spelled range covers the expanded tokens
  72   //  - validating that this range doesn't cover any extra tokens (First/Last)
  73   //
  74   // We do these in order. However as we transform the expanded range into the
  75   // spelled one, we adjust First/Last so the validation remains simple.
  76
  77   assert(SM.getSLocEntry(TargetFile).isFile());
  78   // In most cases, to select First and Last we must return their expansion
  79   // range, i.e. the whole of any macros they are included in.
  80   //
  81   // When First and Last are part of the *same macro arg* of a macro written
  82   // in TargetFile, we that slice of the arg, i.e. their spelling range.
  83   //
  84   // Unwrap such macro calls. If the target file has A(B(C)), the
  85   // SourceLocation stack of a token inside C shows us the expansion of A first,
  86   // then B, then any macros inside C's body, then C itself.
  87   // (This is the reverse of the order the PP applies the expansions in).
  88   while (First.isMacroID() && Last.isMacroID()) {
  89     auto DecFirst = SM.getDecomposedLoc(First);
  90     auto DecLast = SM.getDecomposedLoc(Last);
  91     auto &ExpFirst = SM.getSLocEntry(DecFirst.first).getExpansion();
  92     auto &ExpLast = SM.getSLocEntry(DecLast.first).getExpansion();
  93
  94     if (!ExpFirst.isMacroArgExpansion() || !ExpLast.isMacroArgExpansion())
  95       break;
  96     // Locations are in the same macro arg if they expand to the same place.
  97     // (They may still have different FileIDs - an arg can have >1 chunks!)
  98     if (ExpFirst.getExpansionLocStart() != ExpLast.getExpansionLocStart())
  99       break;
 100     // Careful, given:
 101     //   #define HIDE ID(ID(a))
 102     //   ID(ID(HIDE))
 103     // The token `a` is wrapped in 4 arg-expansions, we only want to unwrap 2.
 104     // We distinguish them by whether the macro expands into the target file.
 105     // Fortunately, the target file ones will always appear first.
 106     auto ExpFileID = SM.getFileID(ExpFirst.getExpansionLocStart());
 107     if (ExpFileID == TargetFile)
 108       break;
 109     // Replace each endpoint with its spelling inside the macro arg.
 110     // (This is getImmediateSpellingLoc without repeating lookups).
 111     First = ExpFirst.getSpellingLoc().getLocWithOffset(DecFirst.second);
 112     Last = ExpLast.getSpellingLoc().getLocWithOffset(DecLast.second);
 113   }
 114
 115   // In all remaining cases we need the full containing macros.
 116   // If this overlaps Prev or Next, then no range is possible.
 117   SourceRange Candidate =
 118       SM.getExpansionRange(SourceRange(First, Last)).getAsRange();
 119   auto DecFirst = SM.getDecomposedExpansionLoc(Candidate.getBegin());
 120   auto DecLast = SM.getDecomposedExpansionLoc(Candidate.getEnd());
 121   // Can end up in the wrong file due to bad input or token-pasting shenanigans.
 122   if (Candidate.isInvalid() || DecFirst.first != TargetFile ||
 123       DecLast.first != TargetFile)
 124     return SourceRange();
 125   // Check bounds, which may still be inside macros.
 126   if (Prev.isValid()) {
 127     auto Dec = SM.getDecomposedLoc(SM.getExpansionRange(Prev).getBegin());
 128     if (Dec.first != DecFirst.first || Dec.second >= DecFirst.second)
 129       return SourceRange();
 130   }
 131   if (Next.isValid()) {
 132     auto Dec = SM.getDecomposedLoc(SM.getExpansionRange(Next).getEnd());
 133     if (Dec.first != DecLast.first || Dec.second <= DecLast.second)
 134       return SourceRange();
 135   }
 136   // Now we know that Candidate is a file range that covers [First, Last]
 137   // without encroaching on {Prev, Next}. Ship it!
 138   return Candidate;
 139 }
 140
 141 } // namespace
 142
 143 syntax::Token::Token(SourceLocation Location, unsigned Length,
 144                      tok::TokenKind Kind)
 145     : Location(Location), Length(Length), Kind(Kind) {
 146   assert(Location.isValid());
 147 }
 148
 149 syntax::Token::Token(const clang::Token &T)
 150     : Token(T.getLocation(), T.getLength(), T.getKind()) {
 151   assert(!T.isAnnotation());
 152 }
 153
 154 llvm::StringRef syntax::Token::text(const SourceManager &SM) const {
 155   bool Invalid = false;
 156   const char *Start = SM.getCharacterData(location(), &Invalid);
 157   assert(!Invalid);
 158   return llvm::StringRef(Start, length());
 159 }
 160
 161 FileRange syntax::Token::range(const SourceManager &SM) const {
 162   assert(location().isFileID() && "must be a spelled token");
 163   FileID File;
 164   unsigned StartOffset;
 165   std::tie(File, StartOffset) = SM.getDecomposedLoc(location());
 166   return FileRange(File, StartOffset, StartOffset + length());
 167 }
 168
 169 FileRange syntax::Token::range(const SourceManager &SM,
 170                                const syntax::Token &First,
 171                                const syntax::Token &Last) {
 172   auto F = First.range(SM);
 173   auto L = Last.range(SM);
 174   assert(F.file() == L.file() && "tokens from different files");
 175   assert((F == L || F.endOffset() <= L.beginOffset()) &&
 176          "wrong order of tokens");
 177   return FileRange(F.file(), F.beginOffset(), L.endOffset());
 178 }
 179
 180 llvm::raw_ostream &syntax::operator<<(llvm::raw_ostream &OS, const Token &T) {
 181   return OS << T.str();
 182 }
 183
 184 FileRange::FileRange(FileID File, unsigned BeginOffset, unsigned EndOffset)
 185     : File(File), Begin(BeginOffset), End(EndOffset) {
 186   assert(File.isValid());
 187   assert(BeginOffset <= EndOffset);
 188 }
 189
 190 FileRange::FileRange(const SourceManager &SM, SourceLocation BeginLoc,
 191                      unsigned Length) {
 192   assert(BeginLoc.isValid());
 193   assert(BeginLoc.isFileID());
 194
 195   std::tie(File, Begin) = SM.getDecomposedLoc(BeginLoc);
 196   End = Begin + Length;
 197 }
 198 FileRange::FileRange(const SourceManager &SM, SourceLocation BeginLoc,
 199                      SourceLocation EndLoc) {
 200   assert(BeginLoc.isValid());
 201   assert(BeginLoc.isFileID());
 202   assert(EndLoc.isValid());
 203   assert(EndLoc.isFileID());
 204   assert(SM.getFileID(BeginLoc) == SM.getFileID(EndLoc));
 205   assert(SM.getFileOffset(BeginLoc) <= SM.getFileOffset(EndLoc));
 206
 207   std::tie(File, Begin) = SM.getDecomposedLoc(BeginLoc);
 208   End = SM.getFileOffset(EndLoc);
 209 }
 210
 211 llvm::raw_ostream &syntax::operator<<(llvm::raw_ostream &OS,
 212                                       const FileRange &R) {
 213   return OS << llvm::formatv("FileRange(file = {0}, offsets = {1}-{2})",
 214                              R.file().getHashValue(), R.beginOffset(),
 215                              R.endOffset());
 216 }
 217
 218 llvm::StringRef FileRange::text(const SourceManager &SM) const {
 219   bool Invalid = false;
 220   StringRef Text = SM.getBufferData(File, &Invalid);
 221   if (Invalid)
 222     return "";
 223   assert(Begin <= Text.size());
 224   assert(End <= Text.size());
 225   return Text.substr(Begin, length());
 226 }
 227
 228 void TokenBuffer::indexExpandedTokens() {
 229   // No-op if the index is already created.
 230   if (!ExpandedTokIndex.empty())
 231     return;
 232   ExpandedTokIndex.reserve(ExpandedTokens.size());
 233   // Index ExpandedTokens for faster lookups by SourceLocation.
 234   for (size_t I = 0, E = ExpandedTokens.size(); I != E; ++I) {
 235     SourceLocation Loc = ExpandedTokens[I].location();
 236     if (Loc.isValid())
 237       ExpandedTokIndex[Loc] = I;
 238   }
 239 }
 240
 241 llvm::ArrayRef<syntax::Token> TokenBuffer::expandedTokens(SourceRange R) const {
 242   if (R.isInvalid())
 243     return {};
 244   if (!ExpandedTokIndex.empty()) {
 245     // Quick lookup if `R` is a token range.
 246     // This is a huge win since majority of the users use ranges provided by an
 247     // AST. Ranges in AST are token ranges from expanded token stream.
 248     const auto B = ExpandedTokIndex.find(R.getBegin());
 249     const auto E = ExpandedTokIndex.find(R.getEnd());
 250     if (B != ExpandedTokIndex.end() && E != ExpandedTokIndex.end()) {
 251       const Token *L = ExpandedTokens.data() + B->getSecond();
 252       // Add 1 to End to make a half-open range.
 253       const Token *R = ExpandedTokens.data() + E->getSecond() + 1;
 254       if (L > R)
 255         return {};
 256       return {L, R};
 257     }
 258   }
 259   // Slow case. Use `isBeforeInTranslationUnit` to binary search for the
 260   // required range.
 261   return getTokensCovering(expandedTokens(), R, *SourceMgr);
 262 }
 263
 264 CharSourceRange FileRange::toCharRange(const SourceManager &SM) const {
 265   return CharSourceRange(
 266       SourceRange(SM.getComposedLoc(File, Begin), SM.getComposedLoc(File, End)),
 267       /*IsTokenRange=*/false);
 268 }
 269
 270 std::pair<const syntax::Token *, const TokenBuffer::Mapping *>
 271 TokenBuffer::spelledForExpandedToken(const syntax::Token *Expanded) const {
 272   assert(Expanded);
 273   assert(ExpandedTokens.data() <= Expanded &&
 274          Expanded < ExpandedTokens.data() + ExpandedTokens.size());
 275
 276   auto FileIt = Files.find(
 277       SourceMgr->getFileID(SourceMgr->getExpansionLoc(Expanded->location())));
 278   assert(FileIt != Files.end() && "no file for an expanded token");
 279
 280   const MarkedFile &File = FileIt->second;
 281
 282   unsigned ExpandedIndex = Expanded - ExpandedTokens.data();
 283   // Find the first mapping that produced tokens after \p Expanded.
 284   auto It = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
 285     return M.BeginExpanded <= ExpandedIndex;
 286   });
 287   // Our token could only be produced by the previous mapping.
 288   if (It == File.Mappings.begin()) {
 289     // No previous mapping, no need to modify offsets.
 290     return {&File.SpelledTokens[ExpandedIndex - File.BeginExpanded],
 291             /*Mapping=*/nullptr};
 292   }
 293   --It; // 'It' now points to last mapping that started before our token.
 294
 295   // Check if the token is part of the mapping.
 296   if (ExpandedIndex < It->EndExpanded)
 297     return {&File.SpelledTokens[It->BeginSpelled], /*Mapping=*/&*It};
 298
 299   // Not part of the mapping, use the index from previous mapping to compute the
 300   // corresponding spelled token.
 301   return {
 302       &File.SpelledTokens[It->EndSpelled + (ExpandedIndex - It->EndExpanded)],
 303       /*Mapping=*/nullptr};
 304 }
 305
 306 const TokenBuffer::Mapping *
 307 TokenBuffer::mappingStartingBeforeSpelled(const MarkedFile &F,
 308                                           const syntax::Token *Spelled) {
 309   assert(F.SpelledTokens.data() <= Spelled);
 310   unsigned SpelledI = Spelled - F.SpelledTokens.data();
 311   assert(SpelledI < F.SpelledTokens.size());
 312
 313   auto It = llvm::partition_point(F.Mappings, [SpelledI](const Mapping &M) {
 314     return M.BeginSpelled <= SpelledI;
 315   });
 316   if (It == F.Mappings.begin())
 317     return nullptr;
 318   --It;
 319   return &*It;
 320 }
 321
 322 llvm::SmallVector<llvm::ArrayRef<syntax::Token>, 1>
 323 TokenBuffer::expandedForSpelled(llvm::ArrayRef<syntax::Token> Spelled) const {
 324   if (Spelled.empty())
 325     return {};
 326   const auto &File = fileForSpelled(Spelled);
 327
 328   auto *FrontMapping = mappingStartingBeforeSpelled(File, &Spelled.front());
 329   unsigned SpelledFrontI = &Spelled.front() - File.SpelledTokens.data();
 330   assert(SpelledFrontI < File.SpelledTokens.size());
 331   unsigned ExpandedBegin;
 332   if (!FrontMapping) {
 333     // No mapping that starts before the first token of Spelled, we don't have
 334     // to modify offsets.
 335     ExpandedBegin = File.BeginExpanded + SpelledFrontI;
 336   } else if (SpelledFrontI < FrontMapping->EndSpelled) {
 337     // This mapping applies to Spelled tokens.
 338     if (SpelledFrontI != FrontMapping->BeginSpelled) {
 339       // Spelled tokens don't cover the entire mapping, returning empty result.
 340       return {}; // FIXME: support macro arguments.
 341     }
 342     // Spelled tokens start at the beginning of this mapping.
 343     ExpandedBegin = FrontMapping->BeginExpanded;
 344   } else {
 345     // Spelled tokens start after the mapping ends (they start in the hole
 346     // between 2 mappings, or between a mapping and end of the file).
 347     ExpandedBegin =
 348         FrontMapping->EndExpanded + (SpelledFrontI - FrontMapping->EndSpelled);
 349   }
 350
 351   auto *BackMapping = mappingStartingBeforeSpelled(File, &Spelled.back());
 352   unsigned SpelledBackI = &Spelled.back() - File.SpelledTokens.data();
 353   unsigned ExpandedEnd;
 354   if (!BackMapping) {
 355     // No mapping that starts before the last token of Spelled, we don't have to
 356     // modify offsets.
 357     ExpandedEnd = File.BeginExpanded + SpelledBackI + 1;
 358   } else if (SpelledBackI < BackMapping->EndSpelled) {
 359     // This mapping applies to Spelled tokens.
 360     if (SpelledBackI + 1 != BackMapping->EndSpelled) {
 361       // Spelled tokens don't cover the entire mapping, returning empty result.
 362       return {}; // FIXME: support macro arguments.
 363     }
 364     ExpandedEnd = BackMapping->EndExpanded;
 365   } else {
 366     // Spelled tokens end after the mapping ends.
 367     ExpandedEnd =
 368         BackMapping->EndExpanded + (SpelledBackI - BackMapping->EndSpelled) + 1;
 369   }
 370
 371   assert(ExpandedBegin < ExpandedTokens.size());
 372   assert(ExpandedEnd < ExpandedTokens.size());
 373   // Avoid returning empty ranges.
 374   if (ExpandedBegin == ExpandedEnd)
 375     return {};
 376   return {llvm::ArrayRef(ExpandedTokens.data() + ExpandedBegin,
 377                          ExpandedTokens.data() + ExpandedEnd)};
 378 }
 379
 380 llvm::ArrayRef<syntax::Token> TokenBuffer::spelledTokens(FileID FID) const {
 381   auto It = Files.find(FID);
 382   assert(It != Files.end());
 383   return It->second.SpelledTokens;
 384 }
 385
 386 const syntax::Token *
 387 TokenBuffer::spelledTokenContaining(SourceLocation Loc) const {
 388   assert(Loc.isFileID());
 389   const auto *Tok = llvm::partition_point(
 390       spelledTokens(SourceMgr->getFileID(Loc)),
 391       [&](const syntax::Token &Tok) { return Tok.endLocation() <= Loc; });
 392   if (!Tok || Loc < Tok->location())
 393     return nullptr;
 394   return Tok;
 395 }
 396
 397 std::string TokenBuffer::Mapping::str() const {
 398   return std::string(
 399       llvm::formatv("spelled tokens: [{0},{1}), expanded tokens: [{2},{3})",
 400                     BeginSpelled, EndSpelled, BeginExpanded, EndExpanded));
 401 }
 402
 403 std::optional<llvm::ArrayRef<syntax::Token>>
 404 TokenBuffer::spelledForExpanded(llvm::ArrayRef<syntax::Token> Expanded) const {
 405   // In cases of invalid code, AST nodes can have source ranges that include
 406   // the `eof` token. As there's no spelling for this token, exclude it from
 407   // the range.
 408   if (!Expanded.empty() && Expanded.back().kind() == tok::eof) {
 409     Expanded = Expanded.drop_back();
 410   }
 411   // Mapping an empty range is ambiguous in case of empty mappings at either end
 412   // of the range, bail out in that case.
 413   if (Expanded.empty())
 414     return std::nullopt;
 415   const syntax::Token *First = &Expanded.front();
 416   const syntax::Token *Last = &Expanded.back();
 417   auto [FirstSpelled, FirstMapping] = spelledForExpandedToken(First);
 418   auto [LastSpelled, LastMapping] = spelledForExpandedToken(Last);
 419
 420   FileID FID = SourceMgr->getFileID(FirstSpelled->location());
 421   // FIXME: Handle multi-file changes by trying to map onto a common root.
 422   if (FID != SourceMgr->getFileID(LastSpelled->location()))
 423     return std::nullopt;
 424
 425   const MarkedFile &File = Files.find(FID)->second;
 426
 427   // If the range is within one macro argument, the result may be only part of a
 428   // Mapping. We must use the general (SourceManager-based) algorithm.
 429   if (FirstMapping && FirstMapping == LastMapping &&
 430       SourceMgr->isMacroArgExpansion(First->location()) &&
 431       SourceMgr->isMacroArgExpansion(Last->location())) {
 432     // We use excluded Prev/Next token for bounds checking.
 433     SourceLocation Prev = (First == &ExpandedTokens.front())
 434                               ? SourceLocation()
 435                               : (First - 1)->location();
 436     SourceLocation Next = (Last == &ExpandedTokens.back())
 437                               ? SourceLocation()
 438                               : (Last + 1)->location();
 439     SourceRange Range = spelledForExpandedSlow(
 440         First->location(), Last->location(), Prev, Next, FID, *SourceMgr);
 441     if (Range.isInvalid())
 442       return std::nullopt;
 443     return getTokensCovering(File.SpelledTokens, Range, *SourceMgr);
 444   }
 445
 446   // Otherwise, use the fast version based on Mappings.
 447   // Do not allow changes that doesn't cover full expansion.
 448   unsigned FirstExpanded = Expanded.begin() - ExpandedTokens.data();
 449   unsigned LastExpanded = Expanded.end() - ExpandedTokens.data();
 450   if (FirstMapping && FirstExpanded != FirstMapping->BeginExpanded)
 451     return std::nullopt;
 452   if (LastMapping && LastMapping->EndExpanded != LastExpanded)
 453     return std::nullopt;
 454   return llvm::ArrayRef(
 455       FirstMapping ? File.SpelledTokens.data() + FirstMapping->BeginSpelled
 456                    : FirstSpelled,
 457       LastMapping ? File.SpelledTokens.data() + LastMapping->EndSpelled
 458                   : LastSpelled + 1);
 459 }
 460
 461 TokenBuffer::Expansion TokenBuffer::makeExpansion(const MarkedFile &F,
 462                                                   const Mapping &M) const {
 463   Expansion E;
 464   E.Spelled = llvm::ArrayRef(F.SpelledTokens.data() + M.BeginSpelled,
 465                              F.SpelledTokens.data() + M.EndSpelled);
 466   E.Expanded = llvm::ArrayRef(ExpandedTokens.data() + M.BeginExpanded,
 467                               ExpandedTokens.data() + M.EndExpanded);
 468   return E;
 469 }
 470
 471 const TokenBuffer::MarkedFile &
 472 TokenBuffer::fileForSpelled(llvm::ArrayRef<syntax::Token> Spelled) const {
 473   assert(!Spelled.empty());
 474   assert(Spelled.front().location().isFileID() && "not a spelled token");
 475   auto FileIt = Files.find(SourceMgr->getFileID(Spelled.front().location()));
 476   assert(FileIt != Files.end() && "file not tracked by token buffer");
 477   const auto &File = FileIt->second;
 478   assert(File.SpelledTokens.data() <= Spelled.data() &&
 479          Spelled.end() <=
 480              (File.SpelledTokens.data() + File.SpelledTokens.size()) &&
 481          "Tokens not in spelled range");
 482 #ifndef NDEBUG
 483   auto T1 = Spelled.back().location();
 484   auto T2 = File.SpelledTokens.back().location();
 485   assert(T1 == T2 || sourceManager().isBeforeInTranslationUnit(T1, T2));
 486 #endif
 487   return File;
 488 }
 489
 490 std::optional<TokenBuffer::Expansion>
 491 TokenBuffer::expansionStartingAt(const syntax::Token *Spelled) const {
 492   assert(Spelled);
 493   const auto &File = fileForSpelled(*Spelled);
 494
 495   unsigned SpelledIndex = Spelled - File.SpelledTokens.data();
 496   auto M = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
 497     return M.BeginSpelled < SpelledIndex;
 498   });
 499   if (M == File.Mappings.end() || M->BeginSpelled != SpelledIndex)
 500     return std::nullopt;
 501   return makeExpansion(File, *M);
 502 }
 503
 504 std::vector<TokenBuffer::Expansion> TokenBuffer::expansionsOverlapping(
 505     llvm::ArrayRef<syntax::Token> Spelled) const {
 506   if (Spelled.empty())
 507     return {};
 508   const auto &File = fileForSpelled(Spelled);
 509
 510   // Find the first overlapping range, and then copy until we stop overlapping.
 511   unsigned SpelledBeginIndex = Spelled.begin() - File.SpelledTokens.data();
 512   unsigned SpelledEndIndex = Spelled.end() - File.SpelledTokens.data();
 513   auto M = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
 514     return M.EndSpelled <= SpelledBeginIndex;
 515   });
 516   std::vector<TokenBuffer::Expansion> Expansions;
 517   for (; M != File.Mappings.end() && M->BeginSpelled < SpelledEndIndex; ++M)
 518     Expansions.push_back(makeExpansion(File, *M));
 519   return Expansions;
 520 }
 521
 522 llvm::ArrayRef<syntax::Token>
 523 syntax::spelledTokensTouching(SourceLocation Loc,
 524                               llvm::ArrayRef<syntax::Token> Tokens) {
 525   assert(Loc.isFileID());
 526
 527   auto *Right = llvm::partition_point(
 528       Tokens, [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
 529   bool AcceptRight = Right != Tokens.end() && Right->location() <= Loc;
 530   bool AcceptLeft =
 531       Right != Tokens.begin() && (Right - 1)->endLocation() >= Loc;
 532   return llvm::ArrayRef(Right - (AcceptLeft ? 1 : 0),
 533                         Right + (AcceptRight ? 1 : 0));
 534 }
 535
 536 llvm::ArrayRef<syntax::Token>
 537 syntax::spelledTokensTouching(SourceLocation Loc,
 538                               const syntax::TokenBuffer &Tokens) {
 539   return spelledTokensTouching(
 540       Loc, Tokens.spelledTokens(Tokens.sourceManager().getFileID(Loc)));
 541 }
 542
 543 const syntax::Token *
 544 syntax::spelledIdentifierTouching(SourceLocation Loc,
 545                                   llvm::ArrayRef<syntax::Token> Tokens) {
 546   for (const syntax::Token &Tok : spelledTokensTouching(Loc, Tokens)) {
 547     if (Tok.kind() == tok::identifier)
 548       return &Tok;
 549   }
 550   return nullptr;
 551 }
 552
 553 const syntax::Token *
 554 syntax::spelledIdentifierTouching(SourceLocation Loc,
 555                                   const syntax::TokenBuffer &Tokens) {
 556   return spelledIdentifierTouching(
 557       Loc, Tokens.spelledTokens(Tokens.sourceManager().getFileID(Loc)));
 558 }
 559
 560 std::vector<const syntax::Token *>
 561 TokenBuffer::macroExpansions(FileID FID) const {
 562   auto FileIt = Files.find(FID);
 563   assert(FileIt != Files.end() && "file not tracked by token buffer");
 564   auto &File = FileIt->second;
 565   std::vector<const syntax::Token *> Expansions;
 566   auto &Spelled = File.SpelledTokens;
 567   for (auto Mapping : File.Mappings) {
 568     const syntax::Token *Token = &Spelled[Mapping.BeginSpelled];
 569     if (Token->kind() == tok::TokenKind::identifier)
 570       Expansions.push_back(Token);
 571   }
 572   return Expansions;
 573 }
 574
 575 std::vector<syntax::Token> syntax::tokenize(const FileRange &FR,
 576                                             const SourceManager &SM,
 577                                             const LangOptions &LO) {
 578   std::vector<syntax::Token> Tokens;
 579   IdentifierTable Identifiers(LO);
 580   auto AddToken = [&](clang::Token T) {
 581     // Fill the proper token kind for keywords, etc.
 582     if (T.getKind() == tok::raw_identifier && !T.needsCleaning() &&
 583         !T.hasUCN()) { // FIXME: support needsCleaning and hasUCN cases.
 584       clang::IdentifierInfo &II = Identifiers.get(T.getRawIdentifier());
 585       T.setIdentifierInfo(&II);
 586       T.setKind(II.getTokenID());
 587     }
 588     Tokens.push_back(syntax::Token(T));
 589   };
 590
 591   auto SrcBuffer = SM.getBufferData(FR.file());
 592   Lexer L(SM.getLocForStartOfFile(FR.file()), LO, SrcBuffer.data(),
 593           SrcBuffer.data() + FR.beginOffset(),
 594           // We can't make BufEnd point to FR.endOffset, as Lexer requires a
 595           // null terminated buffer.
 596           SrcBuffer.data() + SrcBuffer.size());
 597
 598   clang::Token T;
 599   while (!L.LexFromRawLexer(T) && L.getCurrentBufferOffset() < FR.endOffset())
 600     AddToken(T);
 601   // LexFromRawLexer returns true when it parses the last token of the file, add
 602   // it iff it starts within the range we are interested in.
 603   if (SM.getFileOffset(T.getLocation()) < FR.endOffset())
 604     AddToken(T);
 605   return Tokens;
 606 }
 607
 608 std::vector<syntax::Token> syntax::tokenize(FileID FID, const SourceManager &SM,
 609                                             const LangOptions &LO) {
 610   return tokenize(syntax::FileRange(FID, 0, SM.getFileIDSize(FID)), SM, LO);
 611 }
 612
 613 /// Records information reqired to construct mappings for the token buffer that
 614 /// we are collecting.
 615 class TokenCollector::CollectPPExpansions : public PPCallbacks {
 616 public:
 617   CollectPPExpansions(TokenCollector &C) : Collector(&C) {}
 618
 619   /// Disabled instance will stop reporting anything to TokenCollector.
 620   /// This ensures that uses of the preprocessor after TokenCollector::consume()
 621   /// is called do not access the (possibly invalid) collector instance.
 622   void disable() { Collector = nullptr; }
 623
 624   void MacroExpands(const clang::Token &MacroNameTok, const MacroDefinition &MD,
 625                     SourceRange Range, const MacroArgs *Args) override {
 626     if (!Collector)
 627       return;
 628     const auto &SM = Collector->PP.getSourceManager();
 629     // Only record top-level expansions that directly produce expanded tokens.
 630     // This excludes those where:
 631     //   - the macro use is inside a macro body,
 632     //   - the macro appears in an argument to another macro.
 633     // However macro expansion isn't really a tree, it's token rewrite rules,
 634     // so there are other cases, e.g.
 635     //   #define B(X) X
 636     //   #define A 1 + B
 637     //   A(2)
 638     // Both A and B produce expanded tokens, though the macro name 'B' comes
 639     // from an expansion. The best we can do is merge the mappings for both.
 640
 641     // The *last* token of any top-level macro expansion must be in a file.
 642     // (In the example above, see the closing paren of the expansion of B).
 643     if (!Range.getEnd().isFileID())
 644       return;
 645     // If there's a current expansion that encloses this one, this one can't be
 646     // top-level.
 647     if (LastExpansionEnd.isValid() &&
 648         !SM.isBeforeInTranslationUnit(LastExpansionEnd, Range.getEnd()))
 649       return;
 650
 651     // If the macro invocation (B) starts in a macro (A) but ends in a file,
 652     // we'll create a merged mapping for A + B by overwriting the endpoint for
 653     // A's startpoint.
 654     if (!Range.getBegin().isFileID()) {
 655       Range.setBegin(SM.getExpansionLoc(Range.getBegin()));
 656       assert(Collector->Expansions.count(Range.getBegin()) &&
 657              "Overlapping macros should have same expansion location");
 658     }
 659
 660     Collector->Expansions[Range.getBegin()] = Range.getEnd();
 661     LastExpansionEnd = Range.getEnd();
 662   }
 663   // FIXME: handle directives like #pragma, #include, etc.
 664 private:
 665   TokenCollector *Collector;
 666   /// Used to detect recursive macro expansions.
 667   SourceLocation LastExpansionEnd;
 668 };
 669
 670 /// Fills in the TokenBuffer by tracing the run of a preprocessor. The
 671 /// implementation tracks the tokens, macro expansions and directives coming
 672 /// from the preprocessor and:
 673 /// - for each token, figures out if it is a part of an expanded token stream,
 674 ///   spelled token stream or both. Stores the tokens appropriately.
 675 /// - records mappings from the spelled to expanded token ranges, e.g. for macro
 676 ///   expansions.
 677 /// FIXME: also properly record:
 678 ///          - #include directives,
 679 ///          - #pragma, #line and other PP directives,
 680 ///          - skipped pp regions,
 681 ///          - ...
 682
 683 TokenCollector::TokenCollector(Preprocessor &PP) : PP(PP) {
 684   // Collect the expanded token stream during preprocessing.
 685   PP.setTokenWatcher([this](const clang::Token &T) {
 686     if (T.isAnnotation())
 687       return;
 688     DEBUG_WITH_TYPE("collect-tokens", llvm::dbgs()
 689                                           << "Token: "
 690                                           << syntax::Token(T).dumpForTests(
 691                                                  this->PP.getSourceManager())
 692                                           << "\n"
 693
 694     );
 695     Expanded.push_back(syntax::Token(T));
 696   });
 697   // And locations of macro calls, to properly recover boundaries of those in
 698   // case of empty expansions.
 699   auto CB = std::make_unique<CollectPPExpansions>(*this);
 700   this->Collector = CB.get();
 701   PP.addPPCallbacks(std::move(CB));
 702 }
 703
 704 /// Builds mappings and spelled tokens in the TokenBuffer based on the expanded
 705 /// token stream.
 706 class TokenCollector::Builder {
 707 public:
 708   Builder(std::vector<syntax::Token> Expanded, PPExpansions CollectedExpansions,
 709           const SourceManager &SM, const LangOptions &LangOpts)
 710       : Result(SM), CollectedExpansions(std::move(CollectedExpansions)), SM(SM),
 711         LangOpts(LangOpts) {
 712     Result.ExpandedTokens = std::move(Expanded);
 713   }
 714
 715   TokenBuffer build() && {
 716     assert(!Result.ExpandedTokens.empty());
 717     assert(Result.ExpandedTokens.back().kind() == tok::eof);
 718
 719     // Tokenize every file that contributed tokens to the expanded stream.
 720     buildSpelledTokens();
 721
 722     // The expanded token stream consists of runs of tokens that came from
 723     // the same source (a macro expansion, part of a file etc).
 724     // Between these runs are the logical positions of spelled tokens that
 725     // didn't expand to anything.
 726     while (NextExpanded < Result.ExpandedTokens.size() - 1 /* eof */) {
 727       // Create empty mappings for spelled tokens that expanded to nothing here.
 728       // May advance NextSpelled, but NextExpanded is unchanged.
 729       discard();
 730       // Create mapping for a contiguous run of expanded tokens.
 731       // Advances NextExpanded past the run, and NextSpelled accordingly.
 732       unsigned OldPosition = NextExpanded;
 733       advance();
 734       if (NextExpanded == OldPosition)
 735         diagnoseAdvanceFailure();
 736     }
 737     // If any tokens remain in any of the files, they didn't expand to anything.
 738     // Create empty mappings up until the end of the file.
 739     for (const auto &File : Result.Files)
 740       discard(File.first);
 741
 742 #ifndef NDEBUG
 743     for (auto &pair : Result.Files) {
 744       auto &mappings = pair.second.Mappings;
 745       assert(llvm::is_sorted(mappings, [](const TokenBuffer::Mapping &M1,
 746                                           const TokenBuffer::Mapping &M2) {
 747         return M1.BeginSpelled < M2.BeginSpelled &&
 748                M1.EndSpelled < M2.EndSpelled &&
 749                M1.BeginExpanded < M2.BeginExpanded &&
 750                M1.EndExpanded < M2.EndExpanded;
 751       }));
 752     }
 753 #endif
 754
 755     return std::move(Result);
 756   }
 757
 758 private:
 759   // Consume a sequence of spelled tokens that didn't expand to anything.
 760   // In the simplest case, skips spelled tokens until finding one that produced
 761   // the NextExpanded token, and creates an empty mapping for them.
 762   // If Drain is provided, skips remaining tokens from that file instead.
 763   void discard(std::optional<FileID> Drain = std::nullopt) {
 764     SourceLocation Target =
 765         Drain ? SM.getLocForEndOfFile(*Drain)
 766               : SM.getExpansionLoc(
 767                     Result.ExpandedTokens[NextExpanded].location());
 768     FileID File = SM.getFileID(Target);
 769     const auto &SpelledTokens = Result.Files[File].SpelledTokens;
 770     auto &NextSpelled = this->NextSpelled[File];
 771
 772     TokenBuffer::Mapping Mapping;
 773     Mapping.BeginSpelled = NextSpelled;
 774     // When dropping trailing tokens from a file, the empty mapping should
 775     // be positioned within the file's expanded-token range (at the end).
 776     Mapping.BeginExpanded = Mapping.EndExpanded =
 777         Drain ? Result.Files[*Drain].EndExpanded : NextExpanded;
 778     // We may want to split into several adjacent empty mappings.
 779     // FlushMapping() emits the current mapping and starts a new one.
 780     auto FlushMapping = [&, this] {
 781       Mapping.EndSpelled = NextSpelled;
 782       if (Mapping.BeginSpelled != Mapping.EndSpelled)
 783         Result.Files[File].Mappings.push_back(Mapping);
 784       Mapping.BeginSpelled = NextSpelled;
 785     };
 786
 787     while (NextSpelled < SpelledTokens.size() &&
 788            SpelledTokens[NextSpelled].location() < Target) {
 789       // If we know mapping bounds at [NextSpelled, KnownEnd] (macro expansion)
 790       // then we want to partition our (empty) mapping.
 791       //   [Start, NextSpelled) [NextSpelled, KnownEnd] (KnownEnd, Target)
 792       SourceLocation KnownEnd =
 793           CollectedExpansions.lookup(SpelledTokens[NextSpelled].location());
 794       if (KnownEnd.isValid()) {
 795         FlushMapping(); // Emits [Start, NextSpelled)
 796         while (NextSpelled < SpelledTokens.size() &&
 797                SpelledTokens[NextSpelled].location() <= KnownEnd)
 798           ++NextSpelled;
 799         FlushMapping(); // Emits [NextSpelled, KnownEnd]
 800         // Now the loop continues and will emit (KnownEnd, Target).
 801       } else {
 802         ++NextSpelled;
 803       }
 804     }
 805     FlushMapping();
 806   }
 807
 808   // Consumes the NextExpanded token and others that are part of the same run.
 809   // Increases NextExpanded and NextSpelled by at least one, and adds a mapping
 810   // (unless this is a run of file tokens, which we represent with no mapping).
 811   void advance() {
 812     const syntax::Token &Tok = Result.ExpandedTokens[NextExpanded];
 813     SourceLocation Expansion = SM.getExpansionLoc(Tok.location());
 814     FileID File = SM.getFileID(Expansion);
 815     const auto &SpelledTokens = Result.Files[File].SpelledTokens;
 816     auto &NextSpelled = this->NextSpelled[File];
 817
 818     if (Tok.location().isFileID()) {
 819       // A run of file tokens continues while the expanded/spelled tokens match.
 820       while (NextSpelled < SpelledTokens.size() &&
 821              NextExpanded < Result.ExpandedTokens.size() &&
 822              SpelledTokens[NextSpelled].location() ==
 823                  Result.ExpandedTokens[NextExpanded].location()) {
 824         ++NextSpelled;
 825         ++NextExpanded;
 826       }
 827       // We need no mapping for file tokens copied to the expanded stream.
 828     } else {
 829       // We found a new macro expansion. We should have its spelling bounds.
 830       auto End = CollectedExpansions.lookup(Expansion);
 831       assert(End.isValid() && "Macro expansion wasn't captured?");
 832
 833       // Mapping starts here...
 834       TokenBuffer::Mapping Mapping;
 835       Mapping.BeginExpanded = NextExpanded;
 836       Mapping.BeginSpelled = NextSpelled;
 837       // ... consumes spelled tokens within bounds we captured ...
 838       while (NextSpelled < SpelledTokens.size() &&
 839              SpelledTokens[NextSpelled].location() <= End)
 840         ++NextSpelled;
 841       // ... consumes expanded tokens rooted at the same expansion ...
 842       while (NextExpanded < Result.ExpandedTokens.size() &&
 843              SM.getExpansionLoc(
 844                  Result.ExpandedTokens[NextExpanded].location()) == Expansion)
 845         ++NextExpanded;
 846       // ... and ends here.
 847       Mapping.EndExpanded = NextExpanded;
 848       Mapping.EndSpelled = NextSpelled;
 849       Result.Files[File].Mappings.push_back(Mapping);
 850     }
 851   }
 852
 853   // advance() is supposed to consume at least one token - if not, we crash.
 854   void diagnoseAdvanceFailure() {
 855 #ifndef NDEBUG
 856     // Show the failed-to-map token in context.
 857     for (unsigned I = (NextExpanded < 10) ? 0 : NextExpanded - 10;
 858          I < NextExpanded + 5 && I < Result.ExpandedTokens.size(); ++I) {
 859       const char *L =
 860           (I == NextExpanded) ? "!! " : (I < NextExpanded) ? "ok " : "   ";
 861       llvm::errs() << L << Result.ExpandedTokens[I].dumpForTests(SM) << "\n";
 862     }
 863 #endif
 864     llvm_unreachable("Couldn't map expanded token to spelled tokens!");
 865   }
 866
 867   /// Initializes TokenBuffer::Files and fills spelled tokens and expanded
 868   /// ranges for each of the files.
 869   void buildSpelledTokens() {
 870     for (unsigned I = 0; I < Result.ExpandedTokens.size(); ++I) {
 871       const auto &Tok = Result.ExpandedTokens[I];
 872       auto FID = SM.getFileID(SM.getExpansionLoc(Tok.location()));
 873       auto It = Result.Files.try_emplace(FID);
 874       TokenBuffer::MarkedFile &File = It.first->second;
 875
 876       // The eof token should not be considered part of the main-file's range.
 877       File.EndExpanded = Tok.kind() == tok::eof ? I : I + 1;
 878
 879       if (!It.second)
 880         continue; // we have seen this file before.
 881       // This is the first time we see this file.
 882       File.BeginExpanded = I;
 883       File.SpelledTokens = tokenize(FID, SM, LangOpts);
 884     }
 885   }
 886
 887   TokenBuffer Result;
 888   unsigned NextExpanded = 0;                    // cursor in ExpandedTokens
 889   llvm::DenseMap<FileID, unsigned> NextSpelled; // cursor in SpelledTokens
 890   PPExpansions CollectedExpansions;
 891   const SourceManager &SM;
 892   const LangOptions &LangOpts;
 893 };
 894
 895 TokenBuffer TokenCollector::consume() && {
 896   PP.setTokenWatcher(nullptr);
 897   Collector->disable();
 898   return Builder(std::move(Expanded), std::move(Expansions),
 899                  PP.getSourceManager(), PP.getLangOpts())
 900       .build();
 901 }
 902
 903 std::string syntax::Token::str() const {
 904   return std::string(llvm::formatv("Token({0}, length = {1})",
 905                                    tok::getTokenName(kind()), length()));
 906 }
 907
 908 std::string syntax::Token::dumpForTests(const SourceManager &SM) const {
 909   return std::string(llvm::formatv("Token(`{0}`, {1}, length = {2})", text(SM),
 910                                    tok::getTokenName(kind()), length()));
 911 }
 912
 913 std::string TokenBuffer::dumpForTests() const {
 914   auto PrintToken = [this](const syntax::Token &T) -> std::string {
 915     if (T.kind() == tok::eof)
 916       return "<eof>";
 917     return std::string(T.text(*SourceMgr));
 918   };
 919
 920   auto DumpTokens = [this, &PrintToken](llvm::raw_ostream &OS,
 921                                         llvm::ArrayRef<syntax::Token> Tokens) {
 922     if (Tokens.empty()) {
 923       OS << "<empty>";
 924       return;
 925     }
 926     OS << Tokens[0].text(*SourceMgr);
 927     for (unsigned I = 1; I < Tokens.size(); ++I) {
 928       if (Tokens[I].kind() == tok::eof)
 929         continue;
 930       OS << " " << PrintToken(Tokens[I]);
 931     }
 932   };
 933
 934   std::string Dump;
 935   llvm::raw_string_ostream OS(Dump);
 936
 937   OS << "expanded tokens:\n"
 938      << "  ";
 939   // (!) we do not show '<eof>'.
 940   DumpTokens(OS, llvm::ArrayRef(ExpandedTokens).drop_back());
 941   OS << "\n";
 942
 943   std::vector<FileID> Keys;
 944   for (const auto &F : Files)
 945     Keys.push_back(F.first);
 946   llvm::sort(Keys);
 947
 948   for (FileID ID : Keys) {
 949     const MarkedFile &File = Files.find(ID)->second;
 950     auto Entry = SourceMgr->getFileEntryRefForID(ID);
 951     if (!Entry)
 952       continue; // Skip builtin files.
 953     std::string Path = llvm::sys::path::convert_to_slash(Entry->getName());
 954     OS << llvm::formatv("file '{0}'\n", Path) << "  spelled tokens:\n"
 955        << "    ";
 956     DumpTokens(OS, File.SpelledTokens);
 957     OS << "\n";
 958
 959     if (File.Mappings.empty()) {
 960       OS << "  no mappings.\n";
 961       continue;
 962     }
 963     OS << "  mappings:\n";
 964     for (auto &M : File.Mappings) {
 965       OS << llvm::formatv(
 966           "    ['{0}'_{1}, '{2}'_{3}) => ['{4}'_{5}, '{6}'_{7})\n",
 967           PrintToken(File.SpelledTokens[M.BeginSpelled]), M.BeginSpelled,
 968           M.EndSpelled == File.SpelledTokens.size()
 969               ? "<eof>"
 970               : PrintToken(File.SpelledTokens[M.EndSpelled]),
 971           M.EndSpelled, PrintToken(ExpandedTokens[M.BeginExpanded]),
 972           M.BeginExpanded, PrintToken(ExpandedTokens[M.EndExpanded]),
 973           M.EndExpanded);
 974     }
 975   }
 976   return Dump;
 977 }