[clangd] textDocument/documentLink to support include statements with macro argument

Author: Adrien Garbani

clang-tools-extra/clangd/XRefs.cpp

@@ -859,17 +859,56 @@ std::vector<DocumentLink> getDocumentLinks(ParsedAST &AST) {
   for (auto &Inc : AST.getIncludeStructure().MainFileIncludes) {
     if (Inc.Resolved.empty())
       continue;
+
+    // Get the location of the # symbole of the "#include ..." statement
     auto HashLoc = SM.getComposedLoc(SM.getMainFileID(), Inc.HashOffset);
+
+    // get the # Token itself, std::next to get the "include" token and the
+    // first token after (aka "File Token")
     const auto *HashTok = AST.getTokens().spelledTokenContaining(HashLoc);
     assert(HashTok && "got inclusion at wrong offset");
     const auto *IncludeTok = std::next(HashTok);
     const auto *FileTok = std::next(IncludeTok);
-    // FileTok->range is not sufficient here, as raw lexing wouldn't yield
-    // correct tokens for angled filenames. Hence we explicitly use
-    // Inc.Written's length.
-    auto FileRange =
-        syntax::FileRange(SM, FileTok->location(), Inc.Written.length())
-            .toCharRange(SM);
+
+    // The File Token can either be of kind :
+    // "less" if using the "#include <h-char-sequence> new-line" syntax
+    // "string_literal" if using the "#include "q-char-sequence" new-line"
+    // syntax something else (most likely "identifier") if using the "#include
+    // pp-tokens new-line" syntax (#include with macro argument)
+
+    CharSourceRange FileRange;
+
+    if (FileTok->kind() == tok::TokenKind::less) {
+      // FileTok->range would only include the '<' char. Hence we explicitly use
+      // Inc.Written's length.
+      FileRange =
+          syntax::FileRange(SM, FileTok->location(), Inc.Written.length())
+              .toCharRange(SM);
+    } else if (FileTok->kind() == tok::TokenKind::string_literal) {
+      // FileTok->range includes the quotes for string literals so just return
+      // it.
+      FileRange = FileTok->range(SM).toCharRange(SM);
+    } else {
+      // FileTok is the first Token of a macro spelling
+      // We can use the AST to get the macro expansion from the spelling
+      // starting at FileTok and use the expansion to get all the spelled Tokens
+      // that expanded to it
+
+      auto OptExpansion = AST.getTokens().expansionStartingAt(FileTok);
+      if (OptExpansion && !OptExpansion->Spelled.empty()) {
+        // If an expansion was found and has an non-empty spelling, return the
+        // range from the start of the first Token to the end of the last Token
+        const auto &LastTok = OptExpansion->Spelled.back();
+
+        FileRange = FileTok->range(SM).toCharRange(SM);
+        const auto EndRange = LastTok.range(SM).toCharRange(SM);
+        FileRange.setEnd(EndRange.getEnd());
+      } else {
+        // We failed to find a macro expansion from the spelling
+        // fallback to FileTok->range
+        FileRange = FileTok->range(SM).toCharRange(SM);
+      }
+    }
 
     Result.push_back(
         DocumentLink({halfOpenToRange(SM, FileRange),

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

@@ -2776,14 +2776,24 @@ TEST(GetNonLocalDeclRefs, All) {
 
 TEST(DocumentLinks, All) {
   Annotations MainCpp(R"cpp(
+      #define HEADER_AA "faa.h"
+      #define HEADER_BB "fbb.h"
+      #define GET_HEADER(X) HEADER_ ## X 
+
       #/*comments*/include /*comments*/ $foo[["foo.h"]] //more comments
       int end_of_preamble = 0;
       #include $bar[[<bar.h>]]
+      #include $AA[[GET_HEADER(AA)]] // Some comment !
+      # /* What about */ \
+      include /* multiple line */ \
+      $BB[[GET_HEADER( /* statements ? */ \
+      BB /* >>>> Hey ! Dont go further ! this ')' is the last token of the file ! >>>>*/ )]]
     )cpp");
 
   TestTU TU;
   TU.Code = std::string(MainCpp.code());
-  TU.AdditionalFiles = {{"foo.h", ""}, {"bar.h", ""}};
+  TU.AdditionalFiles = {
+      {"faa.h", ""}, {"fbb.h", ""}, {"foo.h", ""}, {"bar.h", ""}};
   TU.ExtraArgs = {"-isystem."};
   auto AST = TU.build();
 
@@ -2793,7 +2803,11 @@ TEST(DocumentLinks, All) {
           DocumentLink({MainCpp.range("foo"),
                         URIForFile::canonicalize(testPath("foo.h"), "")}),
           DocumentLink({MainCpp.range("bar"),
-                        URIForFile::canonicalize(testPath("bar.h"), "")})));
+                        URIForFile::canonicalize(testPath("bar.h"), "")}),
+          DocumentLink({MainCpp.range("AA"),
+                        URIForFile::canonicalize(testPath("faa.h"), "")}),
+          DocumentLink({MainCpp.range("BB"),
+                        URIForFile::canonicalize(testPath("fbb.h"), "")})));
 }
 
 } // namespace