[clangd] Improve inlay hints of things expanded from macros

When we aim a hint at some expanded tokens, we're only willing to attach it
to spelled tokens that exactly corresponde.

e.g.
int zoom(int x, int y, int z);
int dummy = zoom(NUMBERS);

Here we want to place a hint "x:" on the expanded "1", but we shouldn't
be willing to place it on NUMBERS, because it doesn't *exactly*
correspond (it has more tokens).

Fortunately we don't even have to implement this algorithm from scratch,
TokenBuffer has it.

Fixes https://github.com/clangd/clangd/issues/1289
Fixes https://github.com/clangd/clangd/issues/1118
Fixes https://github.com/clangd/clangd/issues/1018

Differential Revision: https://reviews.llvm.org/D133982

(cherry picked from commit 924974a3a13b03090d04860f209ce11b3d9d00a6)

2 files changed, 38 insertions, 30 deletions

diff --git a/clang-tools-extra/clangd/InlayHints.cpp b/clang-tools-extra/clangd/InlayHints.cpp
index 7be05fbc3b34..16c6b1cecc03 100644
--- a/clang-tools-extra/clangd/InlayHints.cpp
+++ b/clang-tools-extra/clangd/InlayHints.cpp
@@ -10,6 +10,7 @@
 #include "Config.h"
 #include "HeuristicResolver.h"
 #include "ParsedAST.h"
+#include "SourceCode.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclarationName.h"
 #include "clang/AST/ExprCXX.h"
@@ -192,8 +193,8 @@ class InlayHintVisitor : public RecursiveASTVisitor<InlayHintVisitor> {
 public:
   InlayHintVisitor(std::vector<InlayHint> &Results, ParsedAST &AST,
                    const Config &Cfg, llvm::Optional<Range> RestrictRange)
-      : Results(Results), AST(AST.getASTContext()), Cfg(Cfg),
-        RestrictRange(std::move(RestrictRange)),
+      : Results(Results), AST(AST.getASTContext()), Tokens(AST.getTokens()),
+        Cfg(Cfg), RestrictRange(std::move(RestrictRange)),
         MainFileID(AST.getSourceManager().getMainFileID()),
         Resolver(AST.getHeuristicResolver()),
         TypeHintPolicy(this->AST.getPrintingPolicy()),
@@ -227,8 +228,7 @@ public:
       return true;
     }
 
-    processCall(E->getParenOrBraceRange().getBegin(), E->getConstructor(),
-                {E->getArgs(), E->getNumArgs()});
+    processCall(E->getConstructor(), {E->getArgs(), E->getNumArgs()});
     return true;
   }
 
@@ -254,7 +254,7 @@ public:
     if (!Callee)
       return true;
 
-    processCall(E->getRParenLoc(), Callee, {E->getArgs(), E->getNumArgs()});
+    processCall(Callee, {E->getArgs(), E->getNumArgs()});
     return true;
   }
 
@@ -278,11 +278,11 @@ public:
     return true;
   }
 
-  void addReturnTypeHint(FunctionDecl *D, SourceLocation Loc) {
+  void addReturnTypeHint(FunctionDecl *D, SourceRange Range) {
     auto *AT = D->getReturnType()->getContainedAutoType();
     if (!AT || AT->getDeducedType().isNull())
       return;
-    addTypeHint(Loc, D->getReturnType(), /*Prefix=*/"-> ");
+    addTypeHint(Range, D->getReturnType(), /*Prefix=*/"-> ");
   }
 
   bool VisitVarDecl(VarDecl *D) {
@@ -375,21 +375,11 @@ public:
 private:
   using NameVec = SmallVector<StringRef, 8>;
 
-  // The purpose of Anchor is to deal with macros. It should be the call's
-  // opening or closing parenthesis or brace. (Always using the opening would
-  // make more sense but CallExpr only exposes the closing.) We heuristically
-  // assume that if this location does not come from a macro definition, then
-  // the entire argument list likely appears in the main file and can be hinted.
-  void processCall(SourceLocation Anchor, const FunctionDecl *Callee,
+  void processCall(const FunctionDecl *Callee,
                    llvm::ArrayRef<const Expr *> Args) {
     if (!Cfg.InlayHints.Parameters || Args.size() == 0 || !Callee)
       return;
 
-    // If the anchor location comes from a macro defintion, there's nowhere to
-    // put hints.
-    if (!AST.getSourceManager().getTopMacroCallerLoc(Anchor).isFileID())
-      return;
-
     // The parameter name of a move or copy constructor is not very interesting.
     if (auto *Ctor = dyn_cast<CXXConstructorDecl>(Callee))
       if (Ctor->isCopyOrMoveConstructor())
@@ -637,25 +627,33 @@ private:
 #undef CHECK_KIND
     }
 
-    auto FileRange =
-        toHalfOpenFileRange(AST.getSourceManager(), AST.getLangOpts(), R);
-    if (!FileRange)
+    auto LSPRange = getHintRange(R);
+    if (!LSPRange)
       return;
-    Range LSPRange{
-        sourceLocToPosition(AST.getSourceManager(), FileRange->getBegin()),
-        sourceLocToPosition(AST.getSourceManager(), FileRange->getEnd())};
-    Position LSPPos = Side == HintSide::Left ? LSPRange.start : LSPRange.end;
+    Position LSPPos = Side == HintSide::Left ? LSPRange->start : LSPRange->end;
     if (RestrictRange &&
         (LSPPos < RestrictRange->start || !(LSPPos < RestrictRange->end)))
       return;
-    // The hint may be in a file other than the main file (for example, a header
-    // file that was included after the preamble), do not show in that case.
-    if (!AST.getSourceManager().isWrittenInMainFile(FileRange->getBegin()))
-      return;
     bool PadLeft = Prefix.consume_front(" ");
     bool PadRight = Suffix.consume_back(" ");
     Results.push_back(InlayHint{LSPPos, (Prefix + Label + Suffix).str(), Kind,
-                                PadLeft, PadRight, LSPRange});
+                                PadLeft, PadRight, *LSPRange});
+  }
+
+  // Get the range of the main file that *exactly* corresponds to R.
+  llvm::Optional<Range> getHintRange(SourceRange R) {
+    const auto &SM = AST.getSourceManager();
+    auto Spelled = Tokens.spelledForExpanded(Tokens.expandedTokens(R));
+    // TokenBuffer will return null if e.g. R corresponds to only part of a
+    // macro expansion.
+    if (!Spelled || Spelled->empty())
+      return llvm::None;
+    // Hint must be within the main file, not e.g. a non-preamble include.
+    if (SM.getFileID(Spelled->front().location()) != SM.getMainFileID() ||
+        SM.getFileID(Spelled->back().location()) != SM.getMainFileID())
+      return llvm::None;
+    return Range{sourceLocToPosition(SM, Spelled->front().location()),
+                 sourceLocToPosition(SM, Spelled->back().endLocation())};
   }
 
   void addTypeHint(SourceRange R, QualType T, llvm::StringRef Prefix) {
@@ -680,6 +678,7 @@ private:
 
   std::vector<InlayHint> &Results;
   ASTContext &AST;
+  const syntax::TokenBuffer &Tokens;
   const Config &Cfg;
   llvm::Optional<Range> RestrictRange;
   FileID MainFileID;
diff --git a/clang-tools-extra/clangd/unittests/InlayHintTests.cpp b/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
index a429c0899187..7127f6cc54c5 100644
--- a/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
+++ b/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
@@ -820,6 +820,15 @@ TEST(ParameterHints, Macros) {
     }
   )cpp",
                        ExpectedHint{"param: ", "param"});
+
+  // If the macro expands to multiple arguments, don't hint it.
+  assertParameterHints(R"cpp(
+    void foo(double x, double y);
+    #define CONSTANTS 3.14, 2.72
+    void bar() {
+      foo(CONSTANTS);
+    }
+  )cpp");
 }
 
 TEST(ParameterHints, ConstructorParens) {