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 *TokenBuffer::spelledTokenAt(SourceLocation Loc) const {
 387   assert(Loc.isFileID());
 388   const auto *Tok = llvm::partition_point(
 389       spelledTokens(SourceMgr->getFileID(Loc)),
 390       [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
 391   if (!Tok || Tok->location() != Loc)
 392     return nullptr;
 393   return Tok;
 394 }
 395
 396 std::string TokenBuffer::Mapping::str() const {
 397   return std::string(
 398       llvm::formatv("spelled tokens: [{0},{1}), expanded tokens: [{2},{3})",
 399                     BeginSpelled, EndSpelled, BeginExpanded, EndExpanded));
 400 }
 401
 402 std::optional<llvm::ArrayRef<syntax::Token>>
 403 TokenBuffer::spelledForExpanded(llvm::ArrayRef<syntax::Token> Expanded) const {
 404   // Mapping an empty range is ambiguous in case of empty mappings at either end
 405   // of the range, bail out in that case.
 406   if (Expanded.empty())
 407     return std::nullopt;
 408   const syntax::Token *First = &Expanded.front();
 409   const syntax::Token *Last = &Expanded.back();
 410   auto [FirstSpelled, FirstMapping] = spelledForExpandedToken(First);
 411   auto [LastSpelled, LastMapping] = spelledForExpandedToken(Last);
 412
 413   FileID FID = SourceMgr->getFileID(FirstSpelled->location());
 414   // FIXME: Handle multi-file changes by trying to map onto a common root.
 415   if (FID != SourceMgr->getFileID(LastSpelled->location()))
 416     return std::nullopt;
 417
 418   const MarkedFile &File = Files.find(FID)->second;
 419
 420   // If the range is within one macro argument, the result may be only part of a
 421   // Mapping. We must use the general (SourceManager-based) algorithm.
 422   if (FirstMapping && FirstMapping == LastMapping &&
 423       SourceMgr->isMacroArgExpansion(First->location()) &&
 424       SourceMgr->isMacroArgExpansion(Last->location())) {
 425     // We use excluded Prev/Next token for bounds checking.
 426     SourceLocation Prev = (First == &ExpandedTokens.front())
 427                               ? SourceLocation()
 428                               : (First - 1)->location();
 429     SourceLocation Next = (Last == &ExpandedTokens.back())
 430                               ? SourceLocation()
 431                               : (Last + 1)->location();
 432     SourceRange Range = spelledForExpandedSlow(
 433         First->location(), Last->location(), Prev, Next, FID, *SourceMgr);
 434     if (Range.isInvalid())
 435       return std::nullopt;
 436     return getTokensCovering(File.SpelledTokens, Range, *SourceMgr);
 437   }
 438
 439   // Otherwise, use the fast version based on Mappings.
 440   // Do not allow changes that doesn't cover full expansion.
 441   unsigned FirstExpanded = Expanded.begin() - ExpandedTokens.data();
 442   unsigned LastExpanded = Expanded.end() - ExpandedTokens.data();
 443   if (FirstMapping && FirstExpanded != FirstMapping->BeginExpanded)
 444     return std::nullopt;
 445   if (LastMapping && LastMapping->EndExpanded != LastExpanded)
 446     return std::nullopt;
 447   return llvm::ArrayRef(
 448       FirstMapping ? File.SpelledTokens.data() + FirstMapping->BeginSpelled
 449                    : FirstSpelled,
 450       LastMapping ? File.SpelledTokens.data() + LastMapping->EndSpelled
 451                   : LastSpelled + 1);
 452 }
 453
 454 TokenBuffer::Expansion TokenBuffer::makeExpansion(const MarkedFile &F,
 455                                                   const Mapping &M) const {
 456   Expansion E;
 457   E.Spelled = llvm::ArrayRef(F.SpelledTokens.data() + M.BeginSpelled,
 458                              F.SpelledTokens.data() + M.EndSpelled);
 459   E.Expanded = llvm::ArrayRef(ExpandedTokens.data() + M.BeginExpanded,
 460                               ExpandedTokens.data() + M.EndExpanded);
 461   return E;
 462 }
 463
 464 const TokenBuffer::MarkedFile &
 465 TokenBuffer::fileForSpelled(llvm::ArrayRef<syntax::Token> Spelled) const {
 466   assert(!Spelled.empty());
 467   assert(Spelled.front().location().isFileID() && "not a spelled token");
 468   auto FileIt = Files.find(SourceMgr->getFileID(Spelled.front().location()));
 469   assert(FileIt != Files.end() && "file not tracked by token buffer");
 470   const auto &File = FileIt->second;
 471   assert(File.SpelledTokens.data() <= Spelled.data() &&
 472          Spelled.end() <=
 473              (File.SpelledTokens.data() + File.SpelledTokens.size()) &&
 474          "Tokens not in spelled range");
 475 #ifndef NDEBUG
 476   auto T1 = Spelled.back().location();
 477   auto T2 = File.SpelledTokens.back().location();
 478   assert(T1 == T2 || sourceManager().isBeforeInTranslationUnit(T1, T2));
 479 #endif
 480   return File;
 481 }
 482
 483 std::optional<TokenBuffer::Expansion>
 484 TokenBuffer::expansionStartingAt(const syntax::Token *Spelled) const {
 485   assert(Spelled);
 486   const auto &File = fileForSpelled(*Spelled);
 487
 488   unsigned SpelledIndex = Spelled - File.SpelledTokens.data();
 489   auto M = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
 490     return M.BeginSpelled < SpelledIndex;
 491   });
 492   if (M == File.Mappings.end() || M->BeginSpelled != SpelledIndex)
 493     return std::nullopt;
 494   return makeExpansion(File, *M);
 495 }
 496
 497 std::vector<TokenBuffer::Expansion> TokenBuffer::expansionsOverlapping(
 498     llvm::ArrayRef<syntax::Token> Spelled) const {
 499   if (Spelled.empty())
 500     return {};
 501   const auto &File = fileForSpelled(Spelled);
 502
 503   // Find the first overlapping range, and then copy until we stop overlapping.
 504   unsigned SpelledBeginIndex = Spelled.begin() - File.SpelledTokens.data();
 505   unsigned SpelledEndIndex = Spelled.end() - File.SpelledTokens.data();
 506   auto M = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
 507     return M.EndSpelled <= SpelledBeginIndex;
 508   });
 509   std::vector<TokenBuffer::Expansion> Expansions;
 510   for (; M != File.Mappings.end() && M->BeginSpelled < SpelledEndIndex; ++M)
 511     Expansions.push_back(makeExpansion(File, *M));
 512   return Expansions;
 513 }
 514
 515 llvm::ArrayRef<syntax::Token>
 516 syntax::spelledTokensTouching(SourceLocation Loc,
 517                               llvm::ArrayRef<syntax::Token> Tokens) {
 518   assert(Loc.isFileID());
 519
 520   auto *Right = llvm::partition_point(
 521       Tokens, [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
 522   bool AcceptRight = Right != Tokens.end() && Right->location() <= Loc;
 523   bool AcceptLeft =
 524       Right != Tokens.begin() && (Right - 1)->endLocation() >= Loc;
 525   return llvm::ArrayRef(Right - (AcceptLeft ? 1 : 0),
 526                         Right + (AcceptRight ? 1 : 0));
 527 }
 528
 529 llvm::ArrayRef<syntax::Token>
 530 syntax::spelledTokensTouching(SourceLocation Loc,
 531                               const syntax::TokenBuffer &Tokens) {
 532   return spelledTokensTouching(
 533       Loc, Tokens.spelledTokens(Tokens.sourceManager().getFileID(Loc)));
 534 }
 535
 536 const syntax::Token *
 537 syntax::spelledIdentifierTouching(SourceLocation Loc,
 538                                   llvm::ArrayRef<syntax::Token> Tokens) {
 539   for (const syntax::Token &Tok : spelledTokensTouching(Loc, Tokens)) {
 540     if (Tok.kind() == tok::identifier)
 541       return &Tok;
 542   }
 543   return nullptr;
 544 }
 545
 546 const syntax::Token *
 547 syntax::spelledIdentifierTouching(SourceLocation Loc,
 548                                   const syntax::TokenBuffer &Tokens) {
 549   return spelledIdentifierTouching(
 550       Loc, Tokens.spelledTokens(Tokens.sourceManager().getFileID(Loc)));
 551 }
 552
 553 std::vector<const syntax::Token *>
 554 TokenBuffer::macroExpansions(FileID FID) const {
 555   auto FileIt = Files.find(FID);
 556   assert(FileIt != Files.end() && "file not tracked by token buffer");
 557   auto &File = FileIt->second;
 558   std::vector<const syntax::Token *> Expansions;
 559   auto &Spelled = File.SpelledTokens;
 560   for (auto Mapping : File.Mappings) {
 561     const syntax::Token *Token = &Spelled[Mapping.BeginSpelled];
 562     if (Token->kind() == tok::TokenKind::identifier)
 563       Expansions.push_back(Token);
 564   }
 565   return Expansions;
 566 }
 567
 568 std::vector<syntax::Token> syntax::tokenize(const FileRange &FR,
 569                                             const SourceManager &SM,
 570                                             const LangOptions &LO) {
 571   std::vector<syntax::Token> Tokens;
 572   IdentifierTable Identifiers(LO);
 573   auto AddToken = [&](clang::Token T) {
 574     // Fill the proper token kind for keywords, etc.
 575     if (T.getKind() == tok::raw_identifier && !T.needsCleaning() &&
 576         !T.hasUCN()) { // FIXME: support needsCleaning and hasUCN cases.
 577       clang::IdentifierInfo &II = Identifiers.get(T.getRawIdentifier());
 578       T.setIdentifierInfo(&II);
 579       T.setKind(II.getTokenID());
 580     }
 581     Tokens.push_back(syntax::Token(T));
 582   };
 583
 584   auto SrcBuffer = SM.getBufferData(FR.file());
 585   Lexer L(SM.getLocForStartOfFile(FR.file()), LO, SrcBuffer.data(),
 586           SrcBuffer.data() + FR.beginOffset(),
 587           // We can't make BufEnd point to FR.endOffset, as Lexer requires a
 588           // null terminated buffer.
 589           SrcBuffer.data() + SrcBuffer.size());
 590
 591   clang::Token T;
 592   while (!L.LexFromRawLexer(T) && L.getCurrentBufferOffset() < FR.endOffset())
 593     AddToken(T);
 594   // LexFromRawLexer returns true when it parses the last token of the file, add
 595   // it iff it starts within the range we are interested in.
 596   if (SM.getFileOffset(T.getLocation()) < FR.endOffset())
 597     AddToken(T);
 598   return Tokens;
 599 }
 600
 601 std::vector<syntax::Token> syntax::tokenize(FileID FID, const SourceManager &SM,
 602                                             const LangOptions &LO) {
 603   return tokenize(syntax::FileRange(FID, 0, SM.getFileIDSize(FID)), SM, LO);
 604 }
 605
 606 /// Records information reqired to construct mappings for the token buffer that
 607 /// we are collecting.
 608 class TokenCollector::CollectPPExpansions : public PPCallbacks {
 609 public:
 610   CollectPPExpansions(TokenCollector &C) : Collector(&C) {}
 611
 612   /// Disabled instance will stop reporting anything to TokenCollector.
 613   /// This ensures that uses of the preprocessor after TokenCollector::consume()
 614   /// is called do not access the (possibly invalid) collector instance.
 615   void disable() { Collector = nullptr; }
 616
 617   void MacroExpands(const clang::Token &MacroNameTok, const MacroDefinition &MD,
 618                     SourceRange Range, const MacroArgs *Args) override {
 619     if (!Collector)
 620       return;
 621     const auto &SM = Collector->PP.getSourceManager();
 622     // Only record top-level expansions that directly produce expanded tokens.
 623     // This excludes those where:
 624     //   - the macro use is inside a macro body,
 625     //   - the macro appears in an argument to another macro.
 626     // However macro expansion isn't really a tree, it's token rewrite rules,
 627     // so there are other cases, e.g.
 628     //   #define B(X) X
 629     //   #define A 1 + B
 630     //   A(2)
 631     // Both A and B produce expanded tokens, though the macro name 'B' comes
 632     // from an expansion. The best we can do is merge the mappings for both.
 633
 634     // The *last* token of any top-level macro expansion must be in a file.
 635     // (In the example above, see the closing paren of the expansion of B).
 636     if (!Range.getEnd().isFileID())
 637       return;
 638     // If there's a current expansion that encloses this one, this one can't be
 639     // top-level.
 640     if (LastExpansionEnd.isValid() &&
 641         !SM.isBeforeInTranslationUnit(LastExpansionEnd, Range.getEnd()))
 642       return;
 643
 644     // If the macro invocation (B) starts in a macro (A) but ends in a file,
 645     // we'll create a merged mapping for A + B by overwriting the endpoint for
 646     // A's startpoint.
 647     if (!Range.getBegin().isFileID()) {
 648       Range.setBegin(SM.getExpansionLoc(Range.getBegin()));
 649       assert(Collector->Expansions.count(Range.getBegin()) &&
 650              "Overlapping macros should have same expansion location");
 651     }
 652
 653     Collector->Expansions[Range.getBegin()] = Range.getEnd();
 654     LastExpansionEnd = Range.getEnd();
 655   }
 656   // FIXME: handle directives like #pragma, #include, etc.
 657 private:
 658   TokenCollector *Collector;
 659   /// Used to detect recursive macro expansions.
 660   SourceLocation LastExpansionEnd;
 661 };
 662
 663 /// Fills in the TokenBuffer by tracing the run of a preprocessor. The
 664 /// implementation tracks the tokens, macro expansions and directives coming
 665 /// from the preprocessor and:
 666 /// - for each token, figures out if it is a part of an expanded token stream,
 667 ///   spelled token stream or both. Stores the tokens appropriately.
 668 /// - records mappings from the spelled to expanded token ranges, e.g. for macro
 669 ///   expansions.
 670 /// FIXME: also properly record:
 671 ///          - #include directives,
 672 ///          - #pragma, #line and other PP directives,
 673 ///          - skipped pp regions,
 674 ///          - ...
 675
 676 TokenCollector::TokenCollector(Preprocessor &PP) : PP(PP) {
 677   // Collect the expanded token stream during preprocessing.
 678   PP.setTokenWatcher([this](const clang::Token &T) {
 679     if (T.isAnnotation())
 680       return;
 681     DEBUG_WITH_TYPE("collect-tokens", llvm::dbgs()
 682                                           << "Token: "
 683                                           << syntax::Token(T).dumpForTests(
 684                                                  this->PP.getSourceManager())
 685                                           << "\n"
 686
 687     );
 688     Expanded.push_back(syntax::Token(T));
 689   });
 690   // And locations of macro calls, to properly recover boundaries of those in
 691   // case of empty expansions.
 692   auto CB = std::make_unique<CollectPPExpansions>(*this);
 693   this->Collector = CB.get();
 694   PP.addPPCallbacks(std::move(CB));
 695 }
 696
 697 /// Builds mappings and spelled tokens in the TokenBuffer based on the expanded
 698 /// token stream.
 699 class TokenCollector::Builder {
 700 public:
 701   Builder(std::vector<syntax::Token> Expanded, PPExpansions CollectedExpansions,
 702           const SourceManager &SM, const LangOptions &LangOpts)
 703       : Result(SM), CollectedExpansions(std::move(CollectedExpansions)), SM(SM),
 704         LangOpts(LangOpts) {
 705     Result.ExpandedTokens = std::move(Expanded);
 706   }
 707
 708   TokenBuffer build() && {
 709     assert(!Result.ExpandedTokens.empty());
 710     assert(Result.ExpandedTokens.back().kind() == tok::eof);
 711
 712     // Tokenize every file that contributed tokens to the expanded stream.
 713     buildSpelledTokens();
 714
 715     // The expanded token stream consists of runs of tokens that came from
 716     // the same source (a macro expansion, part of a file etc).
 717     // Between these runs are the logical positions of spelled tokens that
 718     // didn't expand to anything.
 719     while (NextExpanded < Result.ExpandedTokens.size() - 1 /* eof */) {
 720       // Create empty mappings for spelled tokens that expanded to nothing here.
 721       // May advance NextSpelled, but NextExpanded is unchanged.
 722       discard();
 723       // Create mapping for a contiguous run of expanded tokens.
 724       // Advances NextExpanded past the run, and NextSpelled accordingly.
 725       unsigned OldPosition = NextExpanded;
 726       advance();
 727       if (NextExpanded == OldPosition)
 728         diagnoseAdvanceFailure();
 729     }
 730     // If any tokens remain in any of the files, they didn't expand to anything.
 731     // Create empty mappings up until the end of the file.
 732     for (const auto &File : Result.Files)
 733       discard(File.first);
 734
 735 #ifndef NDEBUG
 736     for (auto &pair : Result.Files) {
 737       auto &mappings = pair.second.Mappings;
 738       assert(llvm::is_sorted(mappings, [](const TokenBuffer::Mapping &M1,
 739                                           const TokenBuffer::Mapping &M2) {
 740         return M1.BeginSpelled < M2.BeginSpelled &&
 741                M1.EndSpelled < M2.EndSpelled &&
 742                M1.BeginExpanded < M2.BeginExpanded &&
 743                M1.EndExpanded < M2.EndExpanded;
 744       }));
 745     }
 746 #endif
 747
 748     return std::move(Result);
 749   }
 750
 751 private:
 752   // Consume a sequence of spelled tokens that didn't expand to anything.
 753   // In the simplest case, skips spelled tokens until finding one that produced
 754   // the NextExpanded token, and creates an empty mapping for them.
 755   // If Drain is provided, skips remaining tokens from that file instead.
 756   void discard(std::optional<FileID> Drain = std::nullopt) {
 757     SourceLocation Target =
 758         Drain ? SM.getLocForEndOfFile(*Drain)
 759               : SM.getExpansionLoc(
 760                     Result.ExpandedTokens[NextExpanded].location());
 761     FileID File = SM.getFileID(Target);
 762     const auto &SpelledTokens = Result.Files[File].SpelledTokens;
 763     auto &NextSpelled = this->NextSpelled[File];
 764
 765     TokenBuffer::Mapping Mapping;
 766     Mapping.BeginSpelled = NextSpelled;
 767     // When dropping trailing tokens from a file, the empty mapping should
 768     // be positioned within the file's expanded-token range (at the end).
 769     Mapping.BeginExpanded = Mapping.EndExpanded =
 770         Drain ? Result.Files[*Drain].EndExpanded : NextExpanded;
 771     // We may want to split into several adjacent empty mappings.
 772     // FlushMapping() emits the current mapping and starts a new one.
 773     auto FlushMapping = [&, this] {
 774       Mapping.EndSpelled = NextSpelled;
 775       if (Mapping.BeginSpelled != Mapping.EndSpelled)
 776         Result.Files[File].Mappings.push_back(Mapping);
 777       Mapping.BeginSpelled = NextSpelled;
 778     };
 779
 780     while (NextSpelled < SpelledTokens.size() &&
 781            SpelledTokens[NextSpelled].location() < Target) {
 782       // If we know mapping bounds at [NextSpelled, KnownEnd] (macro expansion)
 783       // then we want to partition our (empty) mapping.
 784       //   [Start, NextSpelled) [NextSpelled, KnownEnd] (KnownEnd, Target)
 785       SourceLocation KnownEnd =
 786           CollectedExpansions.lookup(SpelledTokens[NextSpelled].location());
 787       if (KnownEnd.isValid()) {
 788         FlushMapping(); // Emits [Start, NextSpelled)
 789         while (NextSpelled < SpelledTokens.size() &&
 790                SpelledTokens[NextSpelled].location() <= KnownEnd)
 791           ++NextSpelled;
 792         FlushMapping(); // Emits [NextSpelled, KnownEnd]
 793         // Now the loop continues and will emit (KnownEnd, Target).
 794       } else {
 795         ++NextSpelled;
 796       }
 797     }
 798     FlushMapping();
 799   }
 800
 801   // Consumes the NextExpanded token and others that are part of the same run.
 802   // Increases NextExpanded and NextSpelled by at least one, and adds a mapping
 803   // (unless this is a run of file tokens, which we represent with no mapping).
 804   void advance() {
 805     const syntax::Token &Tok = Result.ExpandedTokens[NextExpanded];
 806     SourceLocation Expansion = SM.getExpansionLoc(Tok.location());
 807     FileID File = SM.getFileID(Expansion);
 808     const auto &SpelledTokens = Result.Files[File].SpelledTokens;
 809     auto &NextSpelled = this->NextSpelled[File];
 810
 811     if (Tok.location().isFileID()) {
 812       // A run of file tokens continues while the expanded/spelled tokens match.
 813       while (NextSpelled < SpelledTokens.size() &&
 814              NextExpanded < Result.ExpandedTokens.size() &&
 815              SpelledTokens[NextSpelled].location() ==
 816                  Result.ExpandedTokens[NextExpanded].location()) {
 817         ++NextSpelled;
 818         ++NextExpanded;
 819       }
 820       // We need no mapping for file tokens copied to the expanded stream.
 821     } else {
 822       // We found a new macro expansion. We should have its spelling bounds.
 823       auto End = CollectedExpansions.lookup(Expansion);
 824       assert(End.isValid() && "Macro expansion wasn't captured?");
 825
 826       // Mapping starts here...
 827       TokenBuffer::Mapping Mapping;
 828       Mapping.BeginExpanded = NextExpanded;
 829       Mapping.BeginSpelled = NextSpelled;
 830       // ... consumes spelled tokens within bounds we captured ...
 831       while (NextSpelled < SpelledTokens.size() &&
 832              SpelledTokens[NextSpelled].location() <= End)
 833         ++NextSpelled;
 834       // ... consumes expanded tokens rooted at the same expansion ...
 835       while (NextExpanded < Result.ExpandedTokens.size() &&
 836              SM.getExpansionLoc(
 837                  Result.ExpandedTokens[NextExpanded].location()) == Expansion)
 838         ++NextExpanded;
 839       // ... and ends here.
 840       Mapping.EndExpanded = NextExpanded;
 841       Mapping.EndSpelled = NextSpelled;
 842       Result.Files[File].Mappings.push_back(Mapping);
 843     }
 844   }
 845
 846   // advance() is supposed to consume at least one token - if not, we crash.
 847   void diagnoseAdvanceFailure() {
 848 #ifndef NDEBUG
 849     // Show the failed-to-map token in context.
 850     for (unsigned I = (NextExpanded < 10) ? 0 : NextExpanded - 10;
 851          I < NextExpanded + 5 && I < Result.ExpandedTokens.size(); ++I) {
 852       const char *L =
 853           (I == NextExpanded) ? "!! " : (I < NextExpanded) ? "ok " : "   ";
 854       llvm::errs() << L << Result.ExpandedTokens[I].dumpForTests(SM) << "\n";
 855     }
 856 #endif
 857     llvm_unreachable("Couldn't map expanded token to spelled tokens!");
 858   }
 859
 860   /// Initializes TokenBuffer::Files and fills spelled tokens and expanded
 861   /// ranges for each of the files.
 862   void buildSpelledTokens() {
 863     for (unsigned I = 0; I < Result.ExpandedTokens.size(); ++I) {
 864       const auto &Tok = Result.ExpandedTokens[I];
 865       auto FID = SM.getFileID(SM.getExpansionLoc(Tok.location()));
 866       auto It = Result.Files.try_emplace(FID);
 867       TokenBuffer::MarkedFile &File = It.first->second;
 868
 869       // The eof token should not be considered part of the main-file's range.
 870       File.EndExpanded = Tok.kind() == tok::eof ? I : I + 1;
 871
 872       if (!It.second)
 873         continue; // we have seen this file before.
 874       // This is the first time we see this file.
 875       File.BeginExpanded = I;
 876       File.SpelledTokens = tokenize(FID, SM, LangOpts);
 877     }
 878   }
 879
 880   TokenBuffer Result;
 881   unsigned NextExpanded = 0;                    // cursor in ExpandedTokens
 882   llvm::DenseMap<FileID, unsigned> NextSpelled; // cursor in SpelledTokens
 883   PPExpansions CollectedExpansions;
 884   const SourceManager &SM;
 885   const LangOptions &LangOpts;
 886 };
 887
 888 TokenBuffer TokenCollector::consume() && {
 889   PP.setTokenWatcher(nullptr);
 890   Collector->disable();
 891   return Builder(std::move(Expanded), std::move(Expansions),
 892                  PP.getSourceManager(), PP.getLangOpts())
 893       .build();
 894 }
 895
 896 std::string syntax::Token::str() const {
 897   return std::string(llvm::formatv("Token({0}, length = {1})",
 898                                    tok::getTokenName(kind()), length()));
 899 }
 900
 901 std::string syntax::Token::dumpForTests(const SourceManager &SM) const {
 902   return std::string(llvm::formatv("Token(`{0}`, {1}, length = {2})", text(SM),
 903                                    tok::getTokenName(kind()), length()));
 904 }
 905
 906 std::string TokenBuffer::dumpForTests() const {
 907   auto PrintToken = [this](const syntax::Token &T) -> std::string {
 908     if (T.kind() == tok::eof)
 909       return "<eof>";
 910     return std::string(T.text(*SourceMgr));
 911   };
 912
 913   auto DumpTokens = [this, &PrintToken](llvm::raw_ostream &OS,
 914                                         llvm::ArrayRef<syntax::Token> Tokens) {
 915     if (Tokens.empty()) {
 916       OS << "<empty>";
 917       return;
 918     }
 919     OS << Tokens[0].text(*SourceMgr);
 920     for (unsigned I = 1; I < Tokens.size(); ++I) {
 921       if (Tokens[I].kind() == tok::eof)
 922         continue;
 923       OS << " " << PrintToken(Tokens[I]);
 924     }
 925   };
 926
 927   std::string Dump;
 928   llvm::raw_string_ostream OS(Dump);
 929
 930   OS << "expanded tokens:\n"
 931      << "  ";
 932   // (!) we do not show '<eof>'.
 933   DumpTokens(OS, llvm::ArrayRef(ExpandedTokens).drop_back());
 934   OS << "\n";
 935
 936   std::vector<FileID> Keys;
 937   for (const auto &F : Files)
 938     Keys.push_back(F.first);
 939   llvm::sort(Keys);
 940
 941   for (FileID ID : Keys) {
 942     const MarkedFile &File = Files.find(ID)->second;
 943     auto Entry = SourceMgr->getFileEntryRefForID(ID);
 944     if (!Entry)
 945       continue; // Skip builtin files.
 946     std::string Path = llvm::sys::path::convert_to_slash(Entry->getName());
 947     OS << llvm::formatv("file '{0}'\n", Path) << "  spelled tokens:\n"
 948        << "    ";
 949     DumpTokens(OS, File.SpelledTokens);
 950     OS << "\n";
 951
 952     if (File.Mappings.empty()) {
 953       OS << "  no mappings.\n";
 954       continue;
 955     }
 956     OS << "  mappings:\n";
 957     for (auto &M : File.Mappings) {
 958       OS << llvm::formatv(
 959           "    ['{0}'_{1}, '{2}'_{3}) => ['{4}'_{5}, '{6}'_{7})\n",
 960           PrintToken(File.SpelledTokens[M.BeginSpelled]), M.BeginSpelled,
 961           M.EndSpelled == File.SpelledTokens.size()
 962               ? "<eof>"
 963               : PrintToken(File.SpelledTokens[M.EndSpelled]),
 964           M.EndSpelled, PrintToken(ExpandedTokens[M.BeginExpanded]),
 965           M.BeginExpanded, PrintToken(ExpandedTokens[M.EndExpanded]),
 966           M.EndExpanded);
 967     }
 968   }
 969   return Dump;
 970 }