[clang-tools-extra] [clangd] fix extract-to-function for overloaded operators (PR #81640)

[Julian Schmidt via cfe-commits]

https://github.com/5chmidti created https://github.com/llvm/llvm-project/pull/81640

When selecting code that contains the use of overloaded operators,
the SelectionTree will attribute the operator to the operator
declaration, not to the `CXXOperatorCallExpr`. To allow
extract-to-function to work with these operators, make unselected
`CXXOperatorCallExpr`s valid root statements, just like `DeclStmt`s.

Fixes clangd/clangd#1254


>From c9886f24d53de59f11028dc06c56a96ef5896e0e Mon Sep 17 00:00:00 2001
From: Julian Schmidt <git.julian.schmidt at gmail.com>
Date: Tue, 13 Feb 2024 18:59:16 +0100
Subject: [PATCH] [clangd] fix extract-to-function for overloaded operators

When selecting code that contains the use of overloaded operators,
the SelectionTree will attribute the operator to the operator
declaration, not to the `CXXOperatorCallExpr`. To allow
extract-to-function to work with these operators, make unselected
`CXXOperatorCallExpr`s valid root statements, just like `DeclStmt`s.

Fixes clangd/clangd#1254
---
 .../refactor/tweaks/ExtractFunction.cpp       | 15 +++---
 .../unittests/tweaks/ExtractFunctionTests.cpp | 47 +++++++++++++++++++
 clang-tools-extra/docs/ReleaseNotes.rst       |  3 ++
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp b/clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp
index 0302839c58252e..aae480175b33f6 100644
--- a/clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp
+++ b/clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp
@@ -56,6 +56,7 @@
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
+#include "clang/AST/ExprCXX.h"
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Stmt.h"
@@ -70,7 +71,6 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
-#include "llvm/Support/raw_os_ostream.h"
 #include <optional>
 
 namespace clang {
@@ -104,9 +104,12 @@ bool isRootStmt(const Node *N) {
   // Root statement cannot be partially selected.
   if (N->Selected == SelectionTree::Partial)
     return false;
-  // Only DeclStmt can be an unselected RootStmt since VarDecls claim the entire
-  // selection range in selectionTree.
-  if (N->Selected == SelectionTree::Unselected && !N->ASTNode.get<DeclStmt>())
+  // A DeclStmt can be an unselected RootStmt since VarDecls claim the entire
+  // selection range in selectionTree. Additionally, an CXXOperatorCallExpr of a
+  // binary operation can be unselected because it's children claim the entire
+  // selection range in the selection tree (e.g. <<).
+  if (N->Selected == SelectionTree::Unselected && !N->ASTNode.get<DeclStmt>() &&
+      !N->ASTNode.get<CXXOperatorCallExpr>())
     return false;
   return true;
 }
@@ -913,8 +916,8 @@ Expected<Tweak::Effect> ExtractFunction::apply(const Selection &Inputs) {
 
       tooling::Replacements OtherEdit(
           createForwardDeclaration(*ExtractedFunc, SM));
-      if (auto PathAndEdit = Tweak::Effect::fileEdit(SM, SM.getFileID(*FwdLoc),
-                                                 OtherEdit))
+      if (auto PathAndEdit =
+              Tweak::Effect::fileEdit(SM, SM.getFileID(*FwdLoc), OtherEdit))
         MultiFileEffect->ApplyEdits.try_emplace(PathAndEdit->first,
                                                 PathAndEdit->second);
       else
diff --git a/clang-tools-extra/clangd/unittests/tweaks/ExtractFunctionTests.cpp b/clang-tools-extra/clangd/unittests/tweaks/ExtractFunctionTests.cpp
index dec63d454d52c6..8e347b516c6ffe 100644
--- a/clang-tools-extra/clangd/unittests/tweaks/ExtractFunctionTests.cpp
+++ b/clang-tools-extra/clangd/unittests/tweaks/ExtractFunctionTests.cpp
@@ -571,6 +571,53 @@ int getNum(bool Superstitious, int Min, int Max) {
   EXPECT_EQ(apply(Before), After);
 }
 
+TEST_F(ExtractFunctionTest, OverloadedOperators) {
+  Context = File;
+  std::string Before = R"cpp(struct A {
+                int operator+(int x) { return x; }
+              };
+              A &operator<<(A &, int);
+              A &operator|(A &, int);
+
+              A stream{};
+
+              void foo(int, int);
+
+              int main() {
+                [[foo(1, 2);
+                foo(3, 4);
+                stream << 42;
+                stream + 42;
+                stream | 42;
+                foo(1, 2);
+                foo(3, 4);]]
+              })cpp";
+  std::string After =
+      R"cpp(struct A {
+                int operator+(int x) { return x; }
+              };
+              A &operator<<(A &, int);
+              A &operator|(A &, int);
+
+              A stream{};
+
+              void foo(int, int);
+
+              void extracted() {
+foo(1, 2);
+                foo(3, 4);
+                stream << 42;
+                stream + 42;
+                stream | 42;
+                foo(1, 2);
+                foo(3, 4);
+}
+int main() {
+                extracted();
+              })cpp";
+  EXPECT_EQ(apply(Before), After);
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/docs/ReleaseNotes.rst b/clang-tools-extra/docs/ReleaseNotes.rst
index f2fba9aa1450d6..01cd2278c1fb18 100644
--- a/clang-tools-extra/docs/ReleaseNotes.rst
+++ b/clang-tools-extra/docs/ReleaseNotes.rst
@@ -69,6 +69,9 @@ Code completion
 Code actions
 ^^^^^^^^^^^^
 
+- Improved the extract-to-function code action to allow extracting statements
+  with overloaded operators like ``<<`` of ``std::ostream``.
+
 Signature help
 ^^^^^^^^^^^^^^