[clangd] Rethink how SelectionTree deals with macros and #includes.

[sammccall]

Summary

The exclusive-claim model is successful at resolving conflicts over tokens
between parent/child or siblings. However it's not the right model for selecting
AST nodes produced by a macro expansion, which can produce several independent
nodes that are equally associated with the macro invocation.
Additionally, any model that only uses the endpoints in a range can fail when
a macro invocation occurs inside the node.

To address this, we use the existing TokenBuffer in more depth. SourceRanges
are translated into an array of expanded tokens we can iterate over, and in
principle process token by token (in practice, batching related tokens helps).

For regular tokens (and macro-arg expansions) we claim the tokens as before,
but for tokens from macro bodies we merely check whether the macro name was
selected. Thus tokens in macro bodies are selected by selecting the macro name.

The aggregation of Selection is now more principled as we need to be able to
call claim()/peek() an arbitrary number of times.

One side-effect of iterating over the expanded tokens is that (usually) nothing
claims preprocessor tokens like macro names and directives.
Rather than fixing this I just left them unclaimed, and removed support for
determining the selectedness of TUDecl.
(That was originally implemented in 90a5bf92ff97b1, but doesn't seem to be very
important or worth the complexity any longer).

The expandedTokens(SourceLocation) helper could be added locally, but seems to
make sense on TokenBuffer.

[sammccall]

This fixes:

• https://github.com/clangd/clangd/issues/126
• https://github.com/clangd/clangd/issues/202

(I've amended the commit message)

[merge_guards_bot]

Build result: pass - 60204 tests passed, 0 failed and 732 were skipped.
Log files: console-log.txt, CMakeCache.txt

[hokein]

clang-tools-extra/clangd/Selection.cpp

49

I might be missing something, I don't understand why we set Partial if Result != New.

446

what's S for this case?

464

nit: maybe use [early-exist](https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code) here?

496

not sure what's the rational behavior, but I think for the following case, we just have the TUDecl in the selection tree, maybe use the whole macro range?

#define M() 123

int d = M(^); // now we only have the TUDecl in the selection tree.

clang-tools-extra/clangd/unittests/SelectionTests.cpp

139

could we have a testcase to cover the "tokens expanded from another #include file" code path?

clang/lib/Tooling/Syntax/Tokens.cpp

125

nit: for code readability, I'd use llvm::ArrayRef<syntax::Token>::iterator type here.

126

I think the parition_point requires the ExpandedTokens is partitioned, but I didn't found any documentation about this guarantee in the code, would be nice to have this in the comment (probably around the ExpandedTokens).

[sammccall]

Address review comments.

[merge_guards_bot]

Build result: fail - 60299 tests passed, 1 failed and 732 were skipped.

failed: Clangd Unit Tests._/ClangdTests/SelectionTest.PathologicalPreprocessor

Log files: console-log.txt, CMakeCache.txt

[sammccall]

No, this patch is busted, and the tests were too simple go catch it.

The issue is that with no exclusivity check, given {{{ MACRO }}} all the enclosing blocks get to count themselves selected.

We need an exclusivity check on the expanded token stream first. This will yield a sequence of slices of tokens not yet claimed. Then each slice gets partitioned by FileID, each partition gets mapped to spelled tokens and checked as in this patch.

I suspect the exclusivity/claiming at the spelled token level is no longer needed.

clang-tools-extra/clangd/Selection.cpp

49

Added a brief comment and an assertion.
Intuitive explanation: Consider Complete = black, Unselected = white, Partial = grey.
White + White -> white, Black + Black -> black, every other mixture is grey.

(But honestly I wrote up the state transition table first and then extracted the logic by staring at it)

446

S is the function parameter, the implication is that it's the range of the vardecl in the example. Expanded the comment a bit.

496

This is definitely better, but isn't trivial to do, and isn't a very important case.
Added a FIXME rather than clutter/delay this patch with it.

clang/lib/Tooling/Syntax/Tokens.cpp

125

I think that hurts readability on two counts:

• it obscures the actual type: a pointer is concrete and familiar, so it's easier to realize that e.g. > is well-defined here
• it makes it harder to understand return {Begin, End} which relies on the fact that the actual type here is Token*

[sammccall]

Rewrote patch with a better approach (claiming expanded tokens rather than spelled).
Added more tests, including one showing a problem with multiple arg expansion.

[merge_guards_bot]

Build result: pass - 60302 tests passed, 0 failed and 732 were skipped.

Log files: console-log.txt, CMakeCache.txt

[hokein]

Thanks, the patch looks good! please also update the patch description.

clang-tools-extra/clangd/Selection.cpp

81

nit: this comment seems not reflect to the code now.

144

nit: remove the redundant start.

clang-tools-extra/clangd/unittests/SelectionTests.cpp

467

I think cases like below are broken:

#define greater(x, y) x > y? x : y
#define abs(x) x > 0 ? x : -x

Selecting the first element for a macro arg seems good to me.

[thakis]

hover.test is failing on Mac: http://45.33.8.238/mac/3308/step_7.txt

Please take a look, and revert if it takes a while.

[sammccall]

Reverted as 905b002c139f039a32ab9bf1fad63d745d12423f

@thakis any information available about that bot configuration?
I'm not seeing any failures on "official" ones, I can't reproduce locally, and the failures don't make a lot of sense to me (lots of off-by-ones everywhere maybe?)

[thakis]

That particular bot does GN builds, but these test failures repro in the cmake build for me on both macs I tried. (Remember that the official mac bots are on greendragon, not on buildbot -- I'd guess it shows there too.)

[thakis]

Mac is happy now, but it fails to build on Windows: http://45.33.8.238/win/3327/step_4.txt

clang-tools-extra/clangd/Selection.h

SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
                 unsigned Offset);
   // Creates a selection tree for the given range in the main file.
   // The range includes bytes [Start, End).
   // If Start == End, uses the same heuristics as SelectionTree(AST, Start).
   SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
                 unsigned Start, unsigned End);
 
   // Describes to what extent an AST node is covered by the selection.
-  enum Selection {
+  enum Selection : unsigned char {
     // The AST node owns no characters covered by the selection.
     // Note that characters owned by children don't count:
     //   if (x == 0) scream();
     //       ^^^^^^
     // The IfStmt would be Unselected because all the selected characters are
     // associated with its children.
     // (Invisible nodes like ImplicitCastExpr are always unselected).
     Unselected,

clang-tools-extra/clangd/Selection.cpp

 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace {
 using Node = SelectionTree::Node;
 using ast_type_traits::DynTypedNode;
 
+// Sentinel value for the selectedness of a node where we've seen no tokens yet.
+// Once we see a token, treated as fully-selected (so far). If we never see a
+// token, it folds into Unselected. These are never exposed publicly,
+constexpr SelectionTree::Selection NoTokens =
+    static_cast<SelectionTree::Selection>(
+        static_cast<unsigned char>(SelectionTree::Complete + 1));
+
+// Nodes start start with NoTokens, and then use this function to aggregate
+// the selectedness as more tokens are found.
+void update(SelectionTree::Selection &Result, SelectionTree::Selection New) {
+  if (Result == NoTokens)
+    Result = New;
+  else if (Result != New)
+    Result = SelectionTree::Partial;
+}
+
 // Identifies which tokens are selected, and evaluates claims of source ranges
 // by AST nodes. Tokens may be claimed only once: first-come, first-served.
 class SelectedTokens {
 public:
   SelectedTokens(llvm::ArrayRef<syntax::Token> Spelled, const SourceManager &SM,
                  unsigned SelBegin, unsigned SelEnd)
       : SelBegin(SelBegin), SelEnd(SelEnd) {
     // Extract bounds and selected-ness for all tokens spelled in the file.
for (const auto& Tok : Spelled) {
         S.Selected = SelectionTree::Complete;
       else if (S.EndOffset > SelBegin && S.StartOffset < SelEnd)
         S.Selected = SelectionTree::Partial;
       else
         S.Selected = SelectionTree::Unselected;
       S.Claimed = false;
     }
   }
 
   // Associates any tokens overlapping [Begin, End) with an AST node.
   // Tokens that were already claimed by another AST node are not claimed again.
   // Updates Result if the node is selected in the sense of SelectionTree.
   void claim(unsigned Begin, unsigned End, SelectionTree::Selection &Result) {
     assert(Begin <= End);
 
     // Fast-path for missing the selection entirely.
     if (Begin >= SelEnd || End <= SelBegin)
for (auto It = Start; It != Tokens.end() && It->StartOffset < End; ++It) {
           PartialSelection |= (It->Selected == SelectionTree::Partial);
         }
       } else {
         // If the node covers an unselected token, it's not completely selected.
         PartialSelection = true;
       }
     }
 
-    // If some tokens were previously claimed (Result != Unselected), we may
-    // upgrade from Partial->Complete, even if no new tokens were claimed.
-    // Important for [[int a]].
-    if (ClaimedAnyToken || Result) {
-      Result = std::max(Result, PartialSelection ? SelectionTree::Partial
-                                                 : SelectionTree::Complete);
-    }
+    if (ClaimedAnyToken)
+      update(Result, PartialSelection ? SelectionTree::Partial
+
+  // Checks whether the token at Offset is selected, and updates Result.
+  // Does not claim the token.
+  void peek(unsigned Offset, SelectionTree::Selection &Result) const {
+    // Find the token, if it exists.
+    auto I = llvm::partition_point(
+        Tokens, [&](const TokInfo &Tok) { return Tok.EndOffset <= Offset; });
+    if (I == Tokens.end() || I->StartOffset != Offset)
+      return;
+    update(Result, I->Selected);
   }
 
 private:
   struct TokInfo {
     unsigned StartOffset;
     unsigned EndOffset;
     SelectionTree::Selection Selected;
     bool Claimed;
     bool operator<(const TokInfo &Other) const {
       return StartOffset < Other.StartOffset;
     }
   };
   std::vector<TokInfo> Tokens;
   unsigned SelBegin, SelEnd;
 };
 
 // Show the type of a node for debugging.
 void printNodeKind(llvm::raw_ostream &OS, const DynTypedNode &N) {
public:
   static std::deque<Node> collect(ASTContext &AST,
                                   const syntax::TokenBuffer &Tokens,
                                   const PrintingPolicy &PP, unsigned Begin,
                                   unsigned End, FileID File) {
     SelectionVisitor V(AST, Tokens, PP, Begin, End, File);
     V.TraverseAST(AST);
     assert(V.Stack.size() == 1 && "Unpaired push/pop?");
     assert(V.Stack.top() == &V.Nodes.front());
-    // We selected TUDecl if tokens were unclaimed (or the file is empty).
-    SelectionTree::Selection UnclaimedTokens = SelectionTree::Unselected;
-    V.Claimed.claim(Begin, End, UnclaimedTokens);
-    if (UnclaimedTokens || V.Nodes.size() == 1) {
-      StringRef FileContent = AST.getSourceManager().getBufferData(File);
-      // Don't require the trailing newlines to be selected.
-      bool SelectedAll = Begin == 0 && End >= FileContent.rtrim().size();
-      V.Stack.top()->Selected =
-          SelectedAll ? SelectionTree::Complete : SelectionTree::Partial;
-    }
     return std::move(V.Nodes);
   }
 
   // We traverse all "well-behaved" nodes the same way:
   //  - push the node onto the stack
   //  - traverse its children recursively
   //  - pop it from the stack
   //  - hit testing: is intersection(node, selection) - union(children) empty?
private:
 
   SelectionVisitor(ASTContext &AST, const syntax::TokenBuffer &Tokens,
                    const PrintingPolicy &PP, unsigned SelBegin, unsigned SelEnd,
                    FileID SelFile)
       : SM(AST.getSourceManager()), LangOpts(AST.getLangOpts()),
 #ifndef NDEBUG
         PrintPolicy(PP),
 #endif
+        Tokens(Tokens),
         Claimed(Tokens.spelledTokens(SelFile), SM, SelBegin, SelEnd),
         SelFile(SelFile),
         SelBeginTokenStart(SM.getFileOffset(Lexer::GetBeginningOfToken(
             SM.getComposedLoc(SelFile, SelBegin), SM, LangOpts))),
         SelEnd(SelEnd) {
     // Ensure we have a node for the TU decl, regardless of traversal scope.
     Nodes.emplace_back();
     Nodes.back().ASTNode = DynTypedNode::create(*AST.getTranslationUnitDecl());
#endif
   // Pushes a node onto the ancestor stack. Pairs with pop().
   // Performs early hit detection for some nodes (on the earlySourceRange).
   void push(DynTypedNode Node) {
     SourceRange Early = earlySourceRange(Node);
     dlog("{1}push: {0}", printNodeToString(Node, PrintPolicy), indent());
     Nodes.emplace_back();
     Nodes.back().ASTNode = std::move(Node);
     Nodes.back().Parent = Stack.top();
+    Nodes.back().Selected = NoTokens;
     Stack.push(&Nodes.back());
     claimRange(Early, Nodes.back().Selected);
-    // Early hit detection never selects the whole node.
-    if (Nodes.back().Selected)
-      Nodes.back().Selected = SelectionTree::Partial;
   }
 
   // Pops a node off the ancestor stack, and finalizes it. Pairs with push().
   // Performs primary hit detection.
   void pop() {
     Node &N = *Stack.top();
     dlog("{1}pop: {0}", printNodeToString(N.ASTNode, PrintPolicy), indent(-1));
     claimRange(N.ASTNode.getSourceRange(), N.Selected);
+    if (N.Selected == NoTokens)
+      N.Selected = SelectionTree::Unselected;
     if (N.Selected || !N.Children.empty()) {
       // Attach to the tree.
       N.Parent->Children.push_back(&N);
     } else {
       // Neither N any children are selected, it doesn't belong in the tree.
       assert(&N == &Nodes.back());
       Nodes.pop_back();
     }
#endif
 
   // Perform hit-testing of a complete Node against the selection.
   // This runs for every node in the AST, and must be fast in common cases.
   // This is usually called from pop(), so we can take children into account.
   // The existing state of Result is relevant (early/late claims can interact).
   void claimRange(SourceRange S, SelectionTree::Selection &Result) {
     if (!S.isValid())
       return;
-    // toHalfOpenFileRange() allows selection of constructs in macro args. e.g:
-    //   #define LOOP_FOREVER(Body) for(;;) { Body }
-    //   void IncrementLots(int &x) {
-    //     LOOP_FOREVER( ++x; )
-    //   }
-    // Selecting "++x" or "x" will do the right thing.
-    auto Range = toHalfOpenFileRange(SM, LangOpts, S);
-    assert(Range && "We should be able to get the File Range");
-    dlog("{1}claimRange: {0}", Range->printToString(SM), indent());
-    auto B = SM.getDecomposedLoc(Range->getBegin());
-    auto E = SM.getDecomposedLoc(Range->getEnd());
-    // Otherwise, nodes in macro expansions can't be selected.
-    if (B.first != SelFile || E.first != SelFile)
-      return;
-    // Attempt to claim the remaining range. If there's nothing to claim, only
-    // children were selected.
-    Claimed.claim(B.second, E.second, Result);
+
+    // We need to iterate over all the expanded tokens that are part of S.
+    // Consider the macro expansion FLAG(x) -> int x = 0;
+    // Neither S.getBegin() nor S.getEnd() are arg expansions, but x is.
+    // The selection FLAG([[x]]) must partially select the VarDecl.
+    llvm::ArrayRef<syntax::Token> Remaining = Tokens.expandedTokens(S);
+    while (!Remaining.empty()) {
+      // Take consecutive tokens from the same context together for efficiency.
+      FileID FID = SM.getFileID(Remaining.front().location());
+      auto Batch = Remaining.take_while([&](const syntax::Token &T) {
+        return SM.getFileID(T.location()) == FID;
+      });
+      assert(!Batch.empty());
+      Remaining = Remaining.drop_front(Batch.size());
+
+      // There are several possible categories of FileID depending on how the
+      // preprocessor was used to generate these tokens:
+      //   main file, #included file, macro args, macro bodies.
+      // We need to identify the main-file tokens that represent Batch, and
+      // determine whether we want to exclusively claim them. Regular tokens
+      // represent one AST construct, but a macro invocation can represent many.
+      if (FID == SelFile) {
+        // Tokens written directly in the main file.
+        // Claim the token exclusively for this node.
+        Claimed.claim(SM.getFileOffset(Batch.front().location()),
+                      SM.getFileOffset(Batch.back().location()) +
+                          Batch.back().length(),
+                      Result);
+      } else if (Batch.front().location().isFileID()) {
+        // Tokens in another file #included into the main file.
+        // Check if the #include is selected, but don't claim it exclusively.
+        for (SourceLocation Loc = Batch.front().location(); Loc.isValid();
+             Loc = SM.getIncludeLoc(SM.getFileID(Loc))) {
+          if (SM.getFileID(Loc) == SelFile) {
+            Claimed.peek(SM.getFileOffset(Loc), Result);
+            break;
+          }
+        }
+      } else {
+        SourceLocation ArgStart =
+            SM.getTopMacroCallerLoc(Batch.front().location());
+        if (SM.getFileID(ArgStart) == SelFile) {
+          // Tokens that were passed as a macro argument.
+          // Claim the token exclusively for this node.
+          // FIXME: this prevents selecting both occurrences of args used twice.
+          SourceLocation ArgEnd =
+              SM.getTopMacroCallerLoc(Batch.back().location());
+          Claimed.claim(SM.getFileOffset(ArgStart),
+                        SM.getFileOffset(ArgEnd) + Batch.back().length(),
+                        Result);
+        } else {
+          // A non-argument macro expansion.
+          // Check if the macro name is selected, don't claim it exclusively.
+          auto Expansion = SM.getDecomposedExpansionLoc(S.getBegin());
+          if (Expansion.first == SelFile)
+            Claimed.peek(Expansion.second, Result);
+        }
+      }
+    }
     if (Result)
-      dlog("{1}hit selection: {0}",
-           SourceRange(SM.getComposedLoc(B.first, B.second),
-                       SM.getComposedLoc(E.first, E.second))
-               .printToString(SM),
-           indent());
+      dlog("{1}hit selection: {0}", S.printToString(SM), indent());
   }
 
   std::string indent(int Offset = 0) {
     // Cast for signed arithmetic.
     int Amount = int(Stack.size()) + Offset;
     assert(Amount >= 0);
     return std::string(Amount, ' ');
   }
 
   SourceManager &SM;
   const LangOptions &LangOpts;
 #ifndef NDEBUG
   const PrintingPolicy &PrintPolicy;
 #endif
   std::stack<Node *> Stack;
+  const syntax::TokenBuffer &Tokens;
   SelectedTokens Claimed;
   std::deque<Node> Nodes; // Stable pointers as we add more nodes.
   FileID SelFile;
   // If the selection start slices a token in half, the beginning of that token.
   // This is useful for checking whether the end of a token range overlaps
   // the selection: range.end < SelBeginTokenStart is equivalent to
   // range.end + measureToken(range.end) < SelBegin (assuming range.end points
   // to a token), and it saves a lex every time.

clang-tools-extra/clangd/unittests/SelectionTests.cpp

Case Cases[] = {
       {
           R"cpp(
             void foo() { [[if (1^11) { return; } else {^ }]] }
           )cpp",
           "IfStmt",
       },
       {
           R"cpp(
+            int x(int);
+            #define M(foo) x(foo)
+            int a = 42;
+            int b = M([[^a]]);
+          )cpp",
+          "DeclRefExpr",
+      },
+      {
+          R"cpp(
             void foo();
             #define CALL_FUNCTION(X) X()
             void bar() { CALL_FUNCTION([[f^o^o]]); }
           )cpp",
           "DeclRefExpr",
       },
       {
           R"cpp(
const char *Cases[] = {
         }
       )cpp",
       R"cpp(
           template <class T>
           struct unique_ptr {};
           void foo(^$C[[unique_ptr<$C[[unique_ptr<$C[[int]]>]]>]]^ a) {}
       )cpp",
       R"cpp(int a = [[5 >^> 1]];)cpp",
-      R"cpp([[
+      R"cpp(
         #define ECHO(X) X
-        ECHO(EC^HO([[$C[[int]]) EC^HO(a]]));
-      ]])cpp",
+        ECHO(EC^HO($C[[int]]) EC^HO(a));
+      )cpp",
       R"cpp( $C[[^$C[[int]] a^]]; )cpp",
       R"cpp( $C[[^$C[[int]] a = $C[[5]]^]]; )cpp",
   };
   for (const char *C : Cases) {
     Annotations Test(C);
     auto AST = TestTU::withCode(Test.code()).build();
     auto T = makeSelectionTree(C, AST);
 
TEST(SelectionTest, Implicit) {
   EXPECT_EQ(Str, &Str->Parent->Parent->ignoreImplicit())
       << "Didn't unwrap " << nodeKind(&Str->Parent->Parent->ignoreImplicit());
 
   EXPECT_EQ("CXXConstructExpr", nodeKind(&Str->outerImplicit()));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/unittests/TweakTests.cpp

void f(int b = [[1]]) {
       goto label;
       label:
         a = [[1]];
     }
   )cpp";
   EXPECT_UNAVAILABLE(UnavailableCases);
 
   // vector of pairs of input and output strings
-  const std::vector<std::pair<llvm::StringLiteral, llvm::StringLiteral>>
+  const std::vector<std::pair<std::string, std::string>>
       InputOutputs = {
           // extraction from variable declaration/assignment
           {R"cpp(void varDecl() {
                    int a = 5 * (4 + (3 [[- 1)]]);
                  })cpp",
            R"cpp(void varDecl() {
                    auto dummy = (3 - 1); int a = 5 * (4 + dummy);
                  })cpp"},
const std::vector<std::pair<std::string, std::string>>
                    auto dummy = PLUS(1+a); int y = dummy;
                  })cpp"},
           // ensure InsertionPoint isn't inside a macro
           {R"cpp(#define LOOP(x) while (1) {a = x;}
                  void f(int a) {
                    if(1)
                     LOOP(5 + [[3]])
                  })cpp",
-           /*FIXME: It should be extracted like this. SelectionTree needs to be
-             * fixed for macros.
             R"cpp(#define LOOP(x) while (1) {a = x;}
                  void f(int a) {
                    auto dummy = 3; if(1)
                     LOOP(5 + dummy)
-                })cpp"},*/
-           R"cpp(#define LOOP(x) while (1) {a = x;}
-                 void f(int a) {
-                   auto dummy = LOOP(5 + 3); if(1)
-                    dummy
                  })cpp"},
           {R"cpp(#define LOOP(x) do {x;} while(1);
                  void f(int a) {
                    if(1)
                     LOOP(5 + [[3]])
                  })cpp",
            R"cpp(#define LOOP(x) do {x;} while(1);
                  void f(int a) {
void f(const int c) {
   std::string TemplateFailInput = R"cpp(
     template<typename T>
     void f() {
       [[int x;]]
     }
   )cpp";
   EXPECT_EQ(apply(TemplateFailInput), "unavailable");
 
-  // FIXME: This should be extractable after selectionTree works correctly for
-  // macros (currently it doesn't select anything for the following case)
-  std::string MacroFailInput = R"cpp(
+  std::string MacroInput = R"cpp(
     #define F(BODY) void f() { BODY }
     F ([[int x = 0;]])
   )cpp";
-  EXPECT_EQ(apply(MacroFailInput), "unavailable");
+  std::string MacroOutput = R"cpp(
+    #define F(BODY) void f() { BODY }
+    void extracted() {
+int x = 0;
+}
+F (extracted();)
+  )cpp";
+  EXPECT_EQ(apply(MacroInput), MacroOutput);
 
   // Shouldn't crash.
   EXPECT_EQ(apply("void f([[int a]]);"), "unavailable");
   // Don't extract if we select the entire function body (CompoundStmt).
   std::string CompoundFailInput = R"cpp(
     void f() [[{
       int a;
     }]]

clang/include/clang/Tooling/Syntax/Tokens.h

public:
   /// point to one of these tokens.
   /// FIXME: figure out how to handle token splitting, e.g. '>>' can be split
   ///        into two '>' tokens by the parser. However, TokenBuffer currently
   ///        keeps it as a single '>>' token.
   llvm::ArrayRef<syntax::Token> expandedTokens() const {
     return ExpandedTokens;
   }
 
+  /// Returns the subrange of expandedTokens() corresponding to the closed
+  /// token range R.
+  llvm::ArrayRef<syntax::Token> expandedTokens(SourceRange R) const;
+
   /// Find the subrange of spelled tokens that produced the corresponding \p
   /// Expanded tokens.
   ///
   /// EXPECTS: \p Expanded is a subrange of expandedTokens().
   ///
   /// Will fail if the expanded tokens do not correspond to a
   /// sequence of spelled tokens. E.g. for the following example:
   ///

clang/lib/Tooling/Syntax/Tokens.cpp

llvm::StringRef FileRange::text(const SourceManager &SM) const {
   StringRef Text = SM.getBufferData(File, &Invalid);
   if (Invalid)
     return "";
   assert(Begin <= Text.size());
   assert(End <= Text.size());
   return Text.substr(Begin, length());
 }
 
+llvm::ArrayRef<syntax::Token> TokenBuffer::expandedTokens(SourceRange R) const {
+  if (R.isInvalid())
+    return {};
+  const Token *Begin =
+      llvm::partition_point(expandedTokens(), [&](const syntax::Token &T) {
+        return SourceMgr->isBeforeInTranslationUnit(T.location(), R.getBegin());
+      });
+  const Token *End =
+      llvm::partition_point(expandedTokens(), [&](const syntax::Token &T) {
+        return !SourceMgr->isBeforeInTranslationUnit(R.getEnd(), T.location());
+      });
+  if (Begin > End)
+    return {};
+  return {Begin, End};
+}
+
 std::pair<const syntax::Token *, const TokenBuffer::Mapping *>
 TokenBuffer::spelledForExpandedToken(const syntax::Token *Expanded) const {
   assert(Expanded);
   assert(ExpandedTokens.data() <= Expanded &&
          Expanded < ExpandedTokens.data() + ExpandedTokens.size());
 
   auto FileIt = Files.find(
       SourceMgr->getFileID(SourceMgr->getExpansionLoc(Expanded->location())));

clang/unittests/Tooling/Syntax/TokensTest.cpp

 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Support/Annotations.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
+#include "gmock/gmock.h"
 #include <cassert>
 #include <cstdlib>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <memory>
 #include <ostream>
 #include <string>
 
recordTokens(R"cpp(
     ID(1)
     #pragma lalala
     not_mapped
   )cpp");
   EXPECT_THAT(Buffer.spelledForExpanded(findExpanded("not_mapped")),
               ValueIs(SameRange(findSpelled("not_mapped"))));
 }
 
+TEST_F(TokenBufferTest, ExpandedTokensForRange) {
+  recordTokens(R"cpp(
+    #define SIGN(X) X##_washere
+    A SIGN(B) C SIGN(D) E SIGN(F) G
+  )cpp");
+
+  SourceRange R(findExpanded("C").front().location(),
+                findExpanded("F_washere").front().location());
+  // Sanity check: expanded and spelled tokens are stored separately.
+  EXPECT_THAT(Buffer.expandedTokens(R),
+              SameRange(findExpanded("C D_washere E F_washere")));
+  EXPECT_THAT(Buffer.expandedTokens(SourceRange()), testing::IsEmpty());
+}
+
 TEST_F(TokenBufferTest, ExpansionStartingAt) {
   // Object-like macro expansions.
   recordTokens(R"cpp(
     #define FOO 3+4
     int a = FOO 1;
     int b = FOO 2;
   )cpp");