[clangd] Function return type hints: support lambdas, don't duplicate "->"

While here, fix an ugliness:
  auto foo()->auto { return 42; }
This (silly) code gains a "-> int" hint. While correct and useful, it renders as
  auto foo()->int->auto { return 42; }
which is confusing enough to do more harm than good I think.

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

2 files changed, 29 insertions, 8 deletions

diff --git a/clang-tools-extra/clangd/InlayHints.cpp b/clang-tools-extra/clangd/InlayHints.cpp
index 671f9a151d40..2f9624f68af3 100644
--- a/clang-tools-extra/clangd/InlayHints.cpp
+++ b/clang-tools-extra/clangd/InlayHints.cpp
@@ -254,17 +254,30 @@ public:
   }
 
   bool VisitFunctionDecl(FunctionDecl *D) {
-    if (auto *AT = D->getReturnType()->getContainedAutoType()) {
-      QualType Deduced = AT->getDeducedType();
-      if (!Deduced.isNull()) {
-        addTypeHint(D->getFunctionTypeLoc().getRParenLoc(), D->getReturnType(),
-                    /*Prefix=*/"-> ");
-      }
+    if (auto *FPT =
+            llvm::dyn_cast<FunctionProtoType>(D->getType().getTypePtr())) {
+      if (!FPT->hasTrailingReturn())
+        addReturnTypeHint(D, D->getFunctionTypeLoc().getRParenLoc());
     }
+    return true;
+  }
 
+  bool VisitLambdaExpr(LambdaExpr *E) {
+    FunctionDecl *D = E->getCallOperator();
+    if (!E->hasExplicitResultType())
+      addReturnTypeHint(D, E->hasExplicitParameters()
+                               ? D->getFunctionTypeLoc().getRParenLoc()
+                               : E->getIntroducerRange().getEnd());
     return true;
   }
 
+  void addReturnTypeHint(FunctionDecl *D, SourceLocation Loc) {
+    auto *AT = D->getReturnType()->getContainedAutoType();
+    if (!AT || AT->getDeducedType().isNull())
+      return;
+    addTypeHint(Loc, D->getReturnType(), /*Prefix=*/"-> ");
+  }
+
   bool VisitVarDecl(VarDecl *D) {
     // Do not show hints for the aggregate in a structured binding,
     // but show hints for the individual bindings.
diff --git a/clang-tools-extra/clangd/unittests/InlayHintTests.cpp b/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
index 6c3ac0d62e0e..d93a6e68a6ee 100644
--- a/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
+++ b/clang-tools-extra/clangd/unittests/InlayHintTests.cpp
@@ -527,13 +527,18 @@ TEST(TypeHints, Lambda) {
   assertTypeHints(R"cpp(
     void f() {
       int cap = 42;
-      auto $L[[L]] = [cap, $init[[init]] = 1 + 1](int a) { 
+      auto $L[[L]] = [cap, $init[[init]] = 1 + 1](int a$ret[[)]] { 
         return a + cap + init; 
       };
     }
   )cpp",
                   ExpectedHint{": (lambda)", "L"},
-                  ExpectedHint{": int", "init"});
+                  ExpectedHint{": int", "init"}, ExpectedHint{"-> int", "ret"});
+
+  // Lambda return hint shown even if no param list.
+  assertTypeHints("auto $L[[x]] = <:$ret[[:>]]{return 42;};",
+                  ExpectedHint{": (lambda)", "L"},
+                  ExpectedHint{"-> int", "ret"});
 }
 
 // Structured bindings tests.
@@ -616,6 +621,9 @@ TEST(TypeHints, ReturnTypeDeduction) {
     // Do not hint `auto` for trailing return type.
     auto f3() -> int;
 
+    // Do not hint when a trailing return type is specified.
+    auto f4() -> auto* { return "foo"; }
+
     // `auto` conversion operator
     struct A {
       operator auto($retConv[[)]] { return 42; }