[clangd] allow extracting to variable for lambda expressions

Support for extracting lambda expressions, e.g. extracting a lambda from a callexpr (e.g. algorithms/ranges) to a named variable.

Reviewed By: nridge

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

Author: 5chmidti

clang-tools-extra/clangd/refactor/tweaks/ExtractVariable.cpp

@@ -12,8 +12,11 @@
 #include "SourceCode.h"
 #include "refactor/Tweak.h"
 #include "clang/AST/ASTContext.h"
+#include "clang/AST/Decl.h"
+#include "clang/AST/DeclCXX.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
+#include "clang/AST/LambdaCapture.h"
 #include "clang/AST/OperationKinds.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Stmt.h"
@@ -74,10 +77,52 @@ computeReferencedDecls(const clang::Expr *Expr) {
   public:
     std::vector<Decl *> ReferencedDecls;
     bool VisitDeclRefExpr(DeclRefExpr *DeclRef) { // NOLINT
+      // Stop the call operator of lambdas from being marked as a referenced
+      // DeclRefExpr in immediately invoked lambdas.
+      if (const auto *const Method =
+              llvm::dyn_cast<CXXMethodDecl>(DeclRef->getDecl());
+          Method != nullptr && Method->getParent()->isLambda()) {
+        return true;
+      }
       ReferencedDecls.push_back(DeclRef->getDecl());
       return true;
     }
+
+    // Local variables declared inside of the selected lambda cannot go out of
+    // scope. The DeclRefExprs that are important are the variables captured,
+    // the DeclRefExprs inside the initializers of init-capture variables,
+    // variables mentioned in trailing return types, constraints and explicit
+    // defaulted template parameters.
+    bool TraverseLambdaExpr(LambdaExpr *LExpr) {
+      for (const auto &[Capture, Initializer] :
+           llvm::zip(LExpr->captures(), LExpr->capture_inits())) {
+        TraverseLambdaCapture(LExpr, &Capture, Initializer);
+      }
+
+      if (clang::Expr *const RequiresClause =
+              LExpr->getTrailingRequiresClause()) {
+        TraverseStmt(RequiresClause);
+      }
+
+      for (auto *const TemplateParam : LExpr->getExplicitTemplateParameters())
+        TraverseDecl(TemplateParam);
+
+      if (auto *const CallOperator = LExpr->getCallOperator()) {
+        TraverseType(CallOperator->getDeclaredReturnType());
+
+        for (auto *const Param : CallOperator->parameters()) {
+          TraverseParmVarDecl(Param);
+        }
+
+        for (auto *const Attr : CallOperator->attrs()) {
+          TraverseAttr(Attr);
+        }
+      }
+
+      return true;
+    }
   };
+
   FindDeclRefsVisitor Visitor;
   Visitor.TraverseStmt(const_cast<Stmt *>(cast<Stmt>(Expr)));
   return Visitor.ReferencedDecls;
@@ -152,10 +197,16 @@ const clang::Stmt *ExtractionContext::computeInsertionPoint() const {
   auto CanExtractOutside =
       [](const SelectionTree::Node *InsertionPoint) -> bool {
     if (const clang::Stmt *Stmt = InsertionPoint->ASTNode.get<clang::Stmt>()) {
-      // Allow all expressions except LambdaExpr since we don't want to extract
-      // from the captures/default arguments of a lambda
-      if (isa<clang::Expr>(Stmt))
-        return !isa<LambdaExpr>(Stmt);
+      if (isa<clang::Expr>(Stmt)) {
+        // Do not allow extraction from the initializer of a defaulted parameter
+        // to a local variable (e.g. a function-local lambda).
+        if (InsertionPoint->Parent->ASTNode.get<ParmVarDecl>() != nullptr) {
+          return false;
+        }
+
+        return true;
+      }
+
       // We don't yet allow extraction from switch/case stmt as we would need to
       // jump over the switch stmt even if there is a CompoundStmt inside the
       // switch. And there are other Stmts which we don't care about (e.g.
@@ -240,7 +291,7 @@ struct ParsedBinaryOperator {
     SelectedOperands.clear();
 
     if (const BinaryOperator *Op =
-        llvm::dyn_cast_or_null<BinaryOperator>(N.ASTNode.get<Expr>())) {
+            llvm::dyn_cast_or_null<BinaryOperator>(N.ASTNode.get<Expr>())) {
       Kind = Op->getOpcode();
       ExprLoc = Op->getExprLoc();
       SelectedOperands = N.Children;
@@ -255,7 +306,7 @@ struct ParsedBinaryOperator {
       Kind = BinaryOperator::getOverloadedOpcode(Op->getOperator());
       ExprLoc = Op->getExprLoc();
       // Not all children are args, there's also the callee (operator).
-      for (const auto* Child : N.Children) {
+      for (const auto *Child : N.Children) {
         const Expr *E = Child->ASTNode.get<Expr>();
         assert(E && "callee and args should be Exprs!");
         if (E == Op->getArg(0) || E == Op->getArg(1))
@@ -376,15 +427,15 @@ bool childExprIsStmt(const Stmt *Outer, const Expr *Inner) {
   if (llvm::isa<SwitchCase>(Outer))
     return true;
   // Control flow statements use condition etc, but not the body.
-  if (const auto* WS = llvm::dyn_cast<WhileStmt>(Outer))
+  if (const auto *WS = llvm::dyn_cast<WhileStmt>(Outer))
     return Inner == WS->getBody();
-  if (const auto* DS = llvm::dyn_cast<DoStmt>(Outer))
+  if (const auto *DS = llvm::dyn_cast<DoStmt>(Outer))
     return Inner == DS->getBody();
-  if (const auto* FS = llvm::dyn_cast<ForStmt>(Outer))
+  if (const auto *FS = llvm::dyn_cast<ForStmt>(Outer))
     return Inner == FS->getBody();
-  if (const auto* FS = llvm::dyn_cast<CXXForRangeStmt>(Outer))
+  if (const auto *FS = llvm::dyn_cast<CXXForRangeStmt>(Outer))
     return Inner == FS->getBody();
-  if (const auto* IS = llvm::dyn_cast<IfStmt>(Outer))
+  if (const auto *IS = llvm::dyn_cast<IfStmt>(Outer))
     return Inner == IS->getThen() || Inner == IS->getElse();
   // Assume all other cases may be actual expressions.
   // This includes the important case of subexpressions (where Outer is Expr).