[Syntax] expose API for expansions overlapping a spelled token range.

[sammccall]

Summary

This allows efficiently accessing all expansions (without iterating over each
token and searching), and also identifying tokens within a range that are
affected by the preprocessor (which is how clangd will use it).

[ArcsinX]

clang/lib/Tooling/Syntax/Tokens.cpp

252

Could we use const auto & here as in lines 419, 434 for consistency?

431

Will it be useful to have similar API with FileID/MarkedFile parameter?
For example, we already have FileID in SelectionTester::SelectionTester.

[sammccall]

auto

clang/lib/Tooling/Syntax/Tokens.cpp

431

FileID is a bit awkward: then you have to have some way to specify which expansions you're interested in (or get those for the whole file and filter them yourself).

We could add an overload for this, but I think Buf.expansionsAffecting(Buf.spelledTokens(FID)) is pretty acceptable for this case.

(MarkedFile can't be used in an API because it's private)

[ArcsinX]

Thanks. This LGTM, but I think I don't have enough experience to accept revisions and could miss something important. So, may be @kadircet could accept it if its OK for him.

[kadircet]

thanks, mostly LG with a comment about some edge case. just wanna know what you think.

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

281

this sounds more like expansionsAffected rather than affecting ? maybe my mental model requires correction, but it feels like spelled tokens are not affected by expansions, it is the other way around.

maybe even expansionsSpawned or triggered?

this is definitely non-blocking though, i am just bikesheding, comment is explanatory enough in any case.

clang/lib/Tooling/Syntax/Tokens.cpp

254

these two are also asserted in fileForSpelled

257

maybe move that into fileForSpelled too ? (to ensure Spelled is contained within the returned file)

431

this might be inconistent with spelledForExpanded from time to time, e.g:

#define FOO(X) 123
#define BAR

void foo() {
  FOO(BA^R);
}

normally BAR has no expansions, but I think it will get merged into outer macro expansion e.g. 123 coming from FOO(BAR). (whereas spelledForExpanded will treat BAR in isolation)

not sure an important limitation but something to keep in mind.

[sammccall]

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

281

Either mental model works I think. Mine is that expansions are rewrite rules - you start with a stream of tokens, and then each expansion mutates it, and you end up with different ones.

Renamed to expansionsOverlapping which seems more neutral and consistent with expansionStartingAt

clang/lib/Tooling/Syntax/Tokens.cpp

431

I don't understand the consistency you're looking for - AFAICS these are different functions that do different things - spelledForExpanded is much higher level.

Is there something in the name or comment I can add/remove to clarify?

normally BAR has no expansions

Tokens don't really "have expansions" - they're *part of* expansions. Generally expansionsOverlapping() will return the expansions themselves that are overlapping, while spelledForExpanded will include the expanded tokens if the whole expansion is in-range.

The former is a building-block (returning Expansions seems like a clear indicator of that) that will usually require some postfiltering, while the latter is fairly directly usable.

[sammccall]

expansionsAffecting->expansionsOverlapping

[kadircet]

thanks, LGTM! (and loved the choice of overlapping)

clang/lib/Tooling/Syntax/Tokens.cpp

431

Tokens don't really "have expansions" - they're *part of* expansions. Generally expansionsOverlapping() will return the expansions themselves that are overlapping, while spelledForExpanded will include the expanded tokens if the whole expansion is in-range.

yeah expansionsOverlapping() will return the expansions themselves that are overlapping bit made my reasoning a little more solid, i am still thinking with the spelled tokens might expand mindset but I suppose spelled tokens might be part of expansions is a better reasoning.

I don't think any extra comments are necessary, at least in this layer.

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

struct Expansion {
     llvm::ArrayRef<syntax::Token> Spelled;
     llvm::ArrayRef<syntax::Token> Expanded;
   };
   /// If \p Spelled starts a mapping (e.g. if it's a macro name or '#' starting
   /// a preprocessor directive) return the subrange of expanded tokens that the
   /// macro expands to.
   llvm::Optional<Expansion>
   expansionStartingAt(const syntax::Token *Spelled) const;
+  /// Returns all expansions (partially) expanded from the specified tokens.
+  /// This is the expansions whose Spelled range intersects \p Spelled.
+  std::vector<Expansion>
+  expansionsAffecting(llvm::ArrayRef<syntax::Token> Spelled) const;
 
   /// Lexed tokens of a file before preprocessing. E.g. for the following input
   ///     #define DECL(name) int name = 10
   ///     DECL(a);
   /// spelledTokens() returns
   ///    {"#", "define", "DECL", "(", "name", ")", "int", "name", "=", "10",
   ///     "DECL", "(", "a", ")", ";"}
   llvm::ArrayRef<syntax::Token> spelledTokens(FileID FID) const;
private:
   spelledForExpandedToken(const syntax::Token *Expanded) const;
 
   /// Returns a mapping starting before \p Spelled token, or nullptr if no
   /// such mapping exists.
   static const Mapping *
   mappingStartingBeforeSpelled(const MarkedFile &F,
                                const syntax::Token *Spelled);
 
+  /// Convert a private Mapping to a public Expansion.
+  Expansion makeExpansion(const MarkedFile &, const Mapping &) const;
+  /// Returns the file that the Spelled tokens are taken from.
+  /// Asserts that they are non-empty, from a tracked file, and in-bounds.
+  const MarkedFile &fileForSpelled(llvm::ArrayRef<syntax::Token> Spelled) const;
+
   /// Token stream produced after preprocessing, conceputally this captures the
   /// same stream as 'clang -E' (excluding the preprocessor directives like
   /// #file, etc.).
   std::vector<syntax::Token> ExpandedTokens;
   llvm::DenseMap<FileID, MarkedFile> Files;
   // The value is never null, pointer instead of reference to avoid disabling
   // implicit assignment operator.
   const SourceManager *SourceMgr;

clang/lib/Tooling/Syntax/Tokens.cpp

TokenBuffer::mappingStartingBeforeSpelled(const MarkedFile &F,
   --It;
   return &*It;
 }
 
 llvm::SmallVector<llvm::ArrayRef<syntax::Token>, 1>
 TokenBuffer::expandedForSpelled(llvm::ArrayRef<syntax::Token> Spelled) const {
   if (Spelled.empty())
     return {};
-  assert(Spelled.front().location().isFileID());
-
-  auto FID = sourceManager().getFileID(Spelled.front().location());
-  auto It = Files.find(FID);
-  assert(It != Files.end());
-
-  const MarkedFile &File = It->second;
+  const auto &File = fileForSpelled(Spelled);
   // `Spelled` must be a subrange of `File.SpelledTokens`.
   assert(File.SpelledTokens.data() <= Spelled.data());
   assert(&Spelled.back() <=
          File.SpelledTokens.data() + File.SpelledTokens.size());
 #ifndef NDEBUG
   auto T1 = Spelled.back().location();
   auto T2 = File.SpelledTokens.back().location();
   assert(T1 == T2 || sourceManager().isBeforeInTranslationUnit(T1, T2));
 #endif
 
   auto *FrontMapping = mappingStartingBeforeSpelled(File, &Spelled.front());
   unsigned SpelledFrontI = &Spelled.front() - File.SpelledTokens.data();
   assert(SpelledFrontI < File.SpelledTokens.size());
TokenBuffer::spelledForExpanded(llvm::ArrayRef<syntax::Token> Expanded) const {
   // All is good, return the result.
   return llvm::makeArrayRef(
       BeginMapping ? File.SpelledTokens.data() + BeginMapping->BeginSpelled
                    : BeginSpelled,
       LastMapping ? File.SpelledTokens.data() + LastMapping->EndSpelled
                   : LastSpelled + 1);
 }
 
+TokenBuffer::Expansion TokenBuffer::makeExpansion(const MarkedFile &F,
+                                                  const Mapping &M) const {
+  Expansion E;
+  E.Spelled = llvm::makeArrayRef(F.SpelledTokens.data() + M.BeginSpelled,
+                                 F.SpelledTokens.data() + M.EndSpelled);
+  E.Expanded = llvm::makeArrayRef(ExpandedTokens.data() + M.BeginExpanded,
+                                  ExpandedTokens.data() + M.EndExpanded);
+  return E;
+}
+
+const TokenBuffer::MarkedFile &
+TokenBuffer::fileForSpelled(llvm::ArrayRef<syntax::Token> Spelled) const {
+  assert(!Spelled.empty());
+  assert(Spelled.front().location().isFileID() && "not a spelled token");
+  auto FileIt = Files.find(SourceMgr->getFileID(Spelled.front().location()));
+  assert(FileIt != Files.end() && "file not tracked by token buffer");
+  const auto &File = FileIt->second;
+  assert(File.SpelledTokens.data() <= Spelled.data() &&
+         Spelled.end() <=
+             (File.SpelledTokens.data() + File.SpelledTokens.size()) &&
+         "Tokens not in spelled range");
+  return File;
+}
+
 llvm::Optional<TokenBuffer::Expansion>
 TokenBuffer::expansionStartingAt(const syntax::Token *Spelled) const {
   assert(Spelled);
-  assert(Spelled->location().isFileID() && "not a spelled token");
-  auto FileIt = Files.find(SourceMgr->getFileID(Spelled->location()));
-  assert(FileIt != Files.end() && "file not tracked by token buffer");
-
-  auto &File = FileIt->second;
-  assert(File.SpelledTokens.data() <= Spelled &&
-         Spelled < (File.SpelledTokens.data() + File.SpelledTokens.size()));
+  const auto &File = fileForSpelled(*Spelled);
 
   unsigned SpelledIndex = Spelled - File.SpelledTokens.data();
   auto M = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
     return M.BeginSpelled < SpelledIndex;
   });
   if (M == File.Mappings.end() || M->BeginSpelled != SpelledIndex)
     return llvm::None;
+  return makeExpansion(File, *M);
+}
 
-  Expansion E;
-  E.Spelled = llvm::makeArrayRef(File.SpelledTokens.data() + M->BeginSpelled,
-                                 File.SpelledTokens.data() + M->EndSpelled);
-  E.Expanded = llvm::makeArrayRef(ExpandedTokens.data() + M->BeginExpanded,
-                                  ExpandedTokens.data() + M->EndExpanded);
-  return E;
+std::vector<TokenBuffer::Expansion>
+TokenBuffer::expansionsAffecting(llvm::ArrayRef<syntax::Token> Spelled) const {
+  if (Spelled.empty())
+    return {};
+  const auto &File = fileForSpelled(Spelled);
+
+  // Find the first overlapping range, and then copy until we stop overlapping.
+  unsigned SpelledBeginIndex = Spelled.begin() - File.SpelledTokens.data();
+  unsigned SpelledEndIndex = Spelled.end() - File.SpelledTokens.data();
+  auto M = llvm::partition_point(File.Mappings, [&](const Mapping &M) {
+    return M.EndSpelled <= SpelledBeginIndex;
+  });
+  std::vector<TokenBuffer::Expansion> Expansions;
+  for (; M != File.Mappings.end() && M->BeginSpelled < SpelledEndIndex; ++M)
+    Expansions.push_back(makeExpansion(File, *M));
+  return Expansions;
 }
+
 llvm::ArrayRef<syntax::Token>
 syntax::spelledTokensTouching(SourceLocation Loc,
                               llvm::ArrayRef<syntax::Token> Tokens) {
   assert(Loc.isFileID());
 
   auto *Right = llvm::partition_point(
       Tokens, [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
   bool AcceptRight = Right != Tokens.end() && Right->location() <= Loc;

clang/unittests/Tooling/Syntax/TokensTest.cpp

 #include <memory>
 #include <ostream>
 #include <string>
 
 using namespace clang;
 using namespace clang::syntax;
 
 using llvm::ValueIs;
+using ::testing::_;
 using ::testing::AllOf;
 using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::Field;
 using ::testing::IsEmpty;
 using ::testing::Matcher;
 using ::testing::Not;
 using ::testing::Pointee;
TEST_F(TokenBufferTest, ExpandedTokensForRange) {
   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) {
+TEST_F(TokenBufferTest, ExpansionsOverlapping) {
   // Object-like macro expansions.
   recordTokens(R"cpp(
     #define FOO 3+4
     int a = FOO 1;
     int b = FOO 2;
   )cpp");
 
-  llvm::ArrayRef<syntax::Token> Foo1 = findSpelled("FOO 1").drop_back();
+  llvm::ArrayRef<syntax::Token> Foo1 = findSpelled("FOO 1");
   EXPECT_THAT(
       Buffer.expansionStartingAt(Foo1.data()),
-      ValueIs(IsExpansion(SameRange(Foo1),
+      ValueIs(IsExpansion(SameRange(Foo1.drop_back()),
+                          SameRange(findExpanded("3 + 4 1").drop_back()))));
+  EXPECT_THAT(
+      Buffer.expansionsAffecting(Foo1),
+      ElementsAre(IsExpansion(SameRange(Foo1.drop_back()),
                               SameRange(findExpanded("3 + 4 1").drop_back()))));
 
-  llvm::ArrayRef<syntax::Token> Foo2 = findSpelled("FOO 2").drop_back();
+  llvm::ArrayRef<syntax::Token> Foo2 = findSpelled("FOO 2");
   EXPECT_THAT(
       Buffer.expansionStartingAt(Foo2.data()),
-      ValueIs(IsExpansion(SameRange(Foo2),
+      ValueIs(IsExpansion(SameRange(Foo2.drop_back()),
                           SameRange(findExpanded("3 + 4 2").drop_back()))));
+  EXPECT_THAT(
+      Buffer.expansionsAffecting(llvm::makeArrayRef(Foo1.begin(), Foo2.end())),
+      ElementsAre(IsExpansion(SameRange(Foo1.drop_back()), _),
+                  IsExpansion(SameRange(Foo2.drop_back()), _)));
 
   // Function-like macro expansions.
   recordTokens(R"cpp(
     #define ID(X) X
     int a = ID(1+2+3);
     int b = ID(ID(2+3+4));
   )cpp");
 
   llvm::ArrayRef<syntax::Token> ID1 = findSpelled("ID ( 1 + 2 + 3 )");
   EXPECT_THAT(Buffer.expansionStartingAt(&ID1.front()),
               ValueIs(IsExpansion(SameRange(ID1),
                                   SameRange(findExpanded("1 + 2 + 3")))));
   // Only the first spelled token should be found.
   for (const auto &T : ID1.drop_front())
     EXPECT_EQ(Buffer.expansionStartingAt(&T), llvm::None);
 
   llvm::ArrayRef<syntax::Token> ID2 = findSpelled("ID ( ID ( 2 + 3 + 4 ) )");
   EXPECT_THAT(Buffer.expansionStartingAt(&ID2.front()),
               ValueIs(IsExpansion(SameRange(ID2),
                                   SameRange(findExpanded("2 + 3 + 4")))));
   // Only the first spelled token should be found.
   for (const auto &T : ID2.drop_front())
     EXPECT_EQ(Buffer.expansionStartingAt(&T), llvm::None);
 
+  EXPECT_THAT(Buffer.expansionsAffecting(llvm::makeArrayRef(
+                  findSpelled("1 + 2").data(), findSpelled("4").data())),
+              ElementsAre(IsExpansion(SameRange(ID1), _),
+                          IsExpansion(SameRange(ID2), _)));
+
   // PP directives.
   recordTokens(R"cpp(
 #define FOO 1
 int a = FOO;
 #pragma once
 int b = 1;
   )cpp");
 
int b = 1;
   llvm::ArrayRef<syntax::Token> PragmaOnce = findSpelled("# pragma once");
   EXPECT_THAT(
       Buffer.expansionStartingAt(&PragmaOnce.front()),
       ValueIs(IsExpansion(SameRange(PragmaOnce),
                           SameRange(findExpanded("int b").take_front(0)))));
   // Only the first spelled token should be found.
   for (const auto &T : PragmaOnce.drop_front())
     EXPECT_EQ(Buffer.expansionStartingAt(&T), llvm::None);
+
+  EXPECT_THAT(
+      Buffer.expansionsAffecting(findSpelled("FOO ; # pragma")),
+      ElementsAre(IsExpansion(SameRange(findSpelled("FOO ;").drop_back()), _),
+                  IsExpansion(SameRange(PragmaOnce), _)));
 }
 
 TEST_F(TokenBufferTest, TokensToFileRange) {
   addFile("./foo.h", "token_from_header");
   llvm::Annotations Code(R"cpp(
     #define FOO token_from_expansion
     #include "./foo.h"
     $all[[$i[[int]] a = FOO;]]