[clangd] fix extract-to-function for overloaded operators

[5chmidti]

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
CXXOperatorCallExprs valid root statements, just like DeclStmts.

Fixes clangd/clangd#1254

[llvmbot]

@llvm/pr-subscribers-clang-tools-extra

@llvm/pr-subscribers-clangd

Author: Julian Schmidt (5chmidti)

Changes

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
CXXOperatorCallExprs valid root statements, just like DeclStmts.

Fixes clangd/clangd#1254

---

Full diff: https://github.com/llvm/llvm-project/pull/81640.diff

3 Files Affected:

• (modified) clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp (+9-6)
• (modified) clang-tools-extra/clangd/unittests/tweaks/ExtractFunctionTests.cpp (+47)
• (modified) clang-tools-extra/docs/ReleaseNotes.rst (+3)

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
 ^^^^^^^^^^^^^^
 

[5chmidti]

Ping

[5chmidti]

Rebase & Ping

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.