From 4839929bed6e69346ff689ffe88750a90ebc0dfa Mon Sep 17 00:00:00 2001
From: Kadir Cetinkaya <kadircet@google.com>
Date: Thu, 21 Jul 2022 12:43:08 +0200
Subject: [PATCH] [clangd] Make forwarding parameter detection logic resilient

This could crash when our heuristic picks the wrong function. Make sure
there is enough parameters in the candidate to prevent those crashes.

Also special case copy/move constructors to make the heuristic work in
presence of those.

Fixes https://github.com/llvm/llvm-project/issues/56620

Differential Revision: https://reviews.llvm.org/D130260
---
 clang-tools-extra/clangd/AST.cpp                   | 26 +++++++++++--
 .../clangd/unittests/InlayHintTests.cpp            | 45 ++++++++++++++++++++++
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/clang-tools-extra/clangd/AST.cpp b/clang-tools-extra/clangd/AST.cpp
index 44c6f49..0c67c54 100644
--- a/clang-tools-extra/clangd/AST.cpp
+++ b/clang-tools-extra/clangd/AST.cpp
@@ -36,6 +36,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
+#include <iterator>
 #include <string>
 #include <vector>
 
@@ -786,6 +787,9 @@ private:
   // inspects the given callee with the given args to check whether it
   // contains Parameters, and sets Info accordingly.
   void handleCall(FunctionDecl *Callee, typename CallExpr::arg_range Args) {
+    // Skip functions with less parameters, they can't be the target.
+    if (Callee->parameters().size() < Parameters.size())
+      return;
     if (std::any_of(Args.begin(), Args.end(), [](const Expr *E) {
           return dyn_cast<PackExpansionExpr>(E) != nullptr;
         })) {
@@ -822,12 +826,20 @@ private:
   // in the given arguments, if it is there.
   llvm::Optional<size_t> findPack(typename CallExpr::arg_range Args) {
     // find the argument directly referring to the first parameter
-    for (auto It = Args.begin(); It != Args.end(); ++It) {
-      const Expr *Arg = *It;
-      if (const auto *RefArg = unwrapForward(Arg)) {
+    assert(Parameters.size() <= static_cast<size_t>(llvm::size(Args)));
+    for (auto Begin = Args.begin(), End = Args.end() - Parameters.size() + 1;
+         Begin != End; ++Begin) {
+      if (const auto *RefArg = unwrapForward(*Begin)) {
         if (Parameters.front() != RefArg->getDecl())
           continue;
-        return std::distance(Args.begin(), It);
+        // Check that this expands all the way until the last parameter.
+        // It's enough to look at the last parameter, because it isn't possible
+        // to expand without expanding all of them.
+        auto ParamEnd = Begin + Parameters.size() - 1;
+        RefArg = unwrapForward(*ParamEnd);
+        if (!RefArg || Parameters.back() != RefArg->getDecl())
+          continue;
+        return std::distance(Args.begin(), Begin);
       }
     }
     return llvm::None;
@@ -872,6 +884,12 @@ private:
   // std::forward.
   static const DeclRefExpr *unwrapForward(const Expr *E) {
     E = E->IgnoreImplicitAsWritten();
+    // There might be an implicit copy/move constructor call on top of the
+    // forwarded arg.
+    // FIXME: Maybe mark implicit calls in the AST to properly filter here.
+    if (const auto *Const = dyn_cast<CXXConstructExpr>(E))
+      if (Const->getConstructor()->isCopyOrMoveConstructor())
+        E = Const->getArg(0)->IgnoreImplicitAsWritten();
     if (const auto *Call = dyn_cast<CallExpr>(E)) {
       const auto Callee = Call->getBuiltinCallee();
       if (Callee == Builtin::BIforward) {
diff --git a/clang-tools-extra/clangd/unittests/InlayHintTests.cpp b/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
index f574118..87bb0cf 100644
--- a/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
+++ b/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
@@ -1429,6 +1429,51 @@ TEST(InlayHints, RestrictRange) {
               ElementsAre(labelIs(": int"), labelIs(": char")));
 }
 
+TEST(ParameterHints, ArgPacksAndConstructors) {
+  assertParameterHints(
+      R"cpp(
+    struct Foo{ Foo(); Foo(int x); };
+    void foo(Foo a, int b);
+    template <typename... Args>
+    void bar(Args... args) {
+      foo(args...);
+    }
+    template <typename... Args>
+    void baz(Args... args) { foo($param1[[Foo{args...}]], $param2[[1]]); }
+
+    template <typename... Args>
+    void bax(Args... args) { foo($param3[[{args...}]], args...); }
+
+    void foo() {
+      bar($param4[[Foo{}]], $param5[[42]]);
+      bar($param6[[42]], $param7[[42]]);
+      baz($param8[[42]]);
+      bax($param9[[42]]);
+    }
+  )cpp",
+      ExpectedHint{"a: ", "param1"}, ExpectedHint{"b: ", "param2"},
+      ExpectedHint{"a: ", "param3"}, ExpectedHint{"a: ", "param4"},
+      ExpectedHint{"b: ", "param5"}, ExpectedHint{"a: ", "param6"},
+      ExpectedHint{"b: ", "param7"}, ExpectedHint{"x: ", "param8"},
+      ExpectedHint{"b: ", "param9"});
+}
+
+TEST(ParameterHints, DoesntExpandAllArgs) {
+  assertParameterHints(
+      R"cpp(
+    void foo(int x, int y);
+    int id(int a, int b, int c);
+    template <typename... Args>
+    void bar(Args... args) {
+      foo(id($param1[[args]], $param2[[1]], $param3[[args]])...);
+    }
+    void foo() {
+      bar(1, 2); // FIXME: We could have `bar(a: 1, a: 2)` here.
+    }
+  )cpp",
+      ExpectedHint{"a: ", "param1"}, ExpectedHint{"b: ", "param2"},
+      ExpectedHint{"c: ", "param3"});
+}
 // FIXME: Low-hanging fruit where we could omit a type hint:
 //  - auto x = TypeName(...);
 //  - auto x = (TypeName) (...);
-- 
2.7.4