[clangd] fix extract-to-function for overloaded operators (#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.

Partially fixes clangd/clangd#1254

---------

Co-authored-by: Nathan Ridge <[email protected]>

Author: 5chmidti

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 {
@@ -95,18 +95,21 @@ enum FunctionDeclKind {
   OutOfLineDefinition
 };
 
-// A RootStmt is a statement that's fully selected including all it's children
-// and it's parent is unselected.
+// A RootStmt is a statement that's fully selected including all its children
+// and its parent is unselected.
 // Check if a node is a root statement.
 bool isRootStmt(const Node *N) {
   if (!N->ASTNode.get<Stmt>())
     return false;
   // 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, a CXXOperatorCallExpr of a
+  // binary operation can be unselected because its 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

clang-tools-extra/clangd/unittests/tweaks/ExtractFunctionTests.cpp

@@ -190,6 +190,18 @@ F (extracted();)
     }]]
   )cpp";
   EXPECT_EQ(apply(CompoundFailInput), "unavailable");
+
+  ExtraArgs.push_back("-std=c++14");
+  // FIXME: Expressions are currently not extracted
+  EXPECT_EQ(apply(R"cpp(
+                void call() { [[1+1]]; }
+            )cpp"),
+            "unavailable");
+  // FIXME: Single expression statements are currently not extracted
+  EXPECT_EQ(apply(R"cpp(
+                void call() { [[1+1;]] }
+            )cpp"),
+            "unavailable");
 }
 
 TEST_F(ExtractFunctionTest, DifferentHeaderSourceTest) {
@@ -571,6 +583,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

clang-tools-extra/docs/ReleaseNotes.rst

@@ -78,6 +78,9 @@ Code actions
 
 - Added `Swap operands` tweak for certain binary operators.
 
+- Improved the extract-to-function code action to allow extracting statements
+  with overloaded operators like ``<<`` of ``std::ostream``.
+
 Signature help
 ^^^^^^^^^^^^^^