[clangd] Expose SelectionTree to code tweaks, and use it for swap if branches.

Summary:
This reduces the per-check implementation burden and redundant work.
It also makes checks range-aware by default (treating the commonAncestor
as if it were a point selection should be good baseline behavior).

Reviewers: ilya-biryukov

Subscribers: ioeric, MaskRay, jkorous, arphaman, kadircet

Tags: #clang

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

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@352876 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 84 insertions, 124 deletions

diff --git a/clangd/ClangdServer.cpp b/clangd/ClangdServer.cpp
index a677c49a..595e0758 100644
--- a/clangd/ClangdServer.cpp
+++ b/clangd/ClangdServer.cpp
@@ -328,23 +328,28 @@ void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
       "Rename", File, Bind(Action, File.str(), NewName.str(), std::move(CB)));
 }
 
+static llvm::Expected<Tweak::Selection>
+tweakSelection(const Range &Sel, const InputsAndAST &AST) {
+  auto Begin = positionToOffset(AST.Inputs.Contents, Sel.start);
+  if (!Begin)
+    return Begin.takeError();
+  auto End = positionToOffset(AST.Inputs.Contents, Sel.end);
+  if (!End)
+    return End.takeError();
+  return Tweak::Selection(AST.AST, *Begin, *End);
+}
+
 void ClangdServer::enumerateTweaks(PathRef File, Range Sel,
                                    Callback<std::vector<TweakRef>> CB) {
   auto Action = [Sel](decltype(CB) CB, std::string File,
                       Expected<InputsAndAST> InpAST) {
     if (!InpAST)
       return CB(InpAST.takeError());
-
-    auto &AST = InpAST->AST;
-    auto CursorLoc = sourceLocationInMainFile(
-        AST.getASTContext().getSourceManager(), Sel.start);
-    if (!CursorLoc)
-      return CB(CursorLoc.takeError());
-    Tweak::Selection Inputs = {InpAST->Inputs.Contents, InpAST->AST,
-                               *CursorLoc};
-
+    auto Selection = tweakSelection(Sel, *InpAST);
+    if (!Selection)
+      return CB(Selection.takeError());
     std::vector<TweakRef> Res;
-    for (auto &T : prepareTweaks(Inputs))
+    for (auto &T : prepareTweaks(*Selection))
       Res.push_back({T->id(), T->title()});
     CB(std::move(Res));
   };
@@ -359,20 +364,14 @@ void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
                       Expected<InputsAndAST> InpAST) {
     if (!InpAST)
       return CB(InpAST.takeError());
-
-    auto &AST = InpAST->AST;
-    auto CursorLoc = sourceLocationInMainFile(
-        AST.getASTContext().getSourceManager(), Sel.start);
-    if (!CursorLoc)
-      return CB(CursorLoc.takeError());
-    Tweak::Selection Inputs = {InpAST->Inputs.Contents, InpAST->AST,
-                               *CursorLoc};
-
-    auto A = prepareTweak(TweakID, Inputs);
+    auto Selection = tweakSelection(Sel, *InpAST);
+    if (!Selection)
+      return CB(Selection.takeError());
+    auto A = prepareTweak(TweakID, *Selection);
     if (!A)
       return CB(A.takeError());
     // FIXME: run formatter on top of resulting replacements.
-    return CB((*A)->apply(Inputs));
+    return CB((*A)->apply(*Selection));
   };
   WorkScheduler.runWithAST(
       "ApplyTweak", File,
diff --git a/clangd/Selection.cpp b/clangd/Selection.cpp
index 41182a8f..ada79891 100644
--- a/clangd/Selection.cpp
+++ b/clangd/Selection.cpp
@@ -112,15 +112,9 @@ private:
   // An optimization for a common case: nodes outside macro expansions that
   // don't intersect the selection may be recursively skipped.
   bool canSafelySkipNode(SourceRange S) {
-<<<<<<< HEAD
     auto B = SM.getDecomposedLoc(S.getBegin());
     auto E = SM.getDecomposedLoc(S.getEnd());
     if (B.first != SelFile || E.first != SelFile)
-=======
-    auto B = SM.getDecomposedLoc(S.getBegin()),
-         E = SM.getDecomposedLoc(S.getEnd());
-    if (B.first != SM.getMainFileID() || E.first != SM.getMainFileID())
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
       return false;
     return B.second >= SelEnd || E.second < SelBeginTokenStart;
   }
@@ -162,7 +156,6 @@ private:
     //     LOOP_FOREVER( ++x; )
     //   }
     // Selecting "++x" or "x" will do the right thing.
-<<<<<<< HEAD
     auto B = SM.getDecomposedLoc(SM.getTopMacroCallerLoc(S.getBegin()));
     auto E = SM.getDecomposedLoc(SM.getTopMacroCallerLoc(S.getEnd()));
     // Otherwise, nodes in macro expansions can't be selected.
@@ -171,16 +164,6 @@ private:
     // Cheap test: is there any overlap at all between the selection and range?
     // Note that E.second is the *start* of the last token, which is why we
     // compare against the "rounded-down" SelBegin.
-=======
-    auto B = SM.getDecomposedLoc(SM.getTopMacroCallerLoc(S.getBegin())),
-         E = SM.getDecomposedLoc(SM.getTopMacroCallerLoc(S.getEnd()));
-    // Otherwise, nodes in macro expansions can't be selected.
-    if (B.first != SM.getMainFileID() || E.first != SM.getMainFileID())
-      return SelectionTree::Unselected;
-    // Cheap test: is there any overlap at all between the selection and range?
-    // Note that E.second is the *start* of the last token, which is why we
-    // compare against the "rounded-down" MinOffset.
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
     if (B.second >= SelEnd || E.second < SelBeginTokenStart)
       return SelectionTree::Unselected;
 
@@ -213,11 +196,7 @@ private:
       CharSourceRange R = SM.getExpansionRange(N->ASTNode.getSourceRange());
       auto B = SM.getDecomposedLoc(R.getBegin());
       auto E = SM.getDecomposedLoc(R.getEnd());
-<<<<<<< HEAD
       if (B.first != SelFile || E.first != SelFile)
-=======
-      if (B.first != SM.getMainFileID() || E.first != SM.getMainFileID())
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
         continue;
       assert(R.isTokenRange());
       // Try to cover up to the next token, spaces between children don't count.
@@ -243,7 +222,6 @@ private:
   SourceManager &SM;
   const LangOptions &LangOpts;
   std::stack<Node *> Stack;
-<<<<<<< HEAD
   std::deque<Node> Nodes; // Stable pointers as we add more nodes.
   // Half-open selection range.
   unsigned SelBegin;
@@ -255,10 +233,6 @@ private:
   // range.end + measureToken(range.end) < SelBegin (assuming range.end points
   // to a token), and it saves a lex every time.
   unsigned SelBeginTokenStart;
-=======
-  std::deque<Node> Nodes;
-  unsigned SelBegin, SelEnd, SelBeginTokenStart;
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
 };
 
 } // namespace
@@ -278,16 +252,9 @@ void SelectionTree::print(llvm::raw_ostream &OS, const SelectionTree::Node &N,
 }
 
 // Decide which selection emulates a "point" query in between characters.
-<<<<<<< HEAD
 static std::pair<unsigned, unsigned> pointBounds(unsigned Offset, FileID FID,
                                                  ASTContext &AST) {
   StringRef Buf = AST.getSourceManager().getBufferData(FID);
-=======
-static std::pair<unsigned, unsigned> pointBounds(unsigned Offset,
-                                                 ASTContext &AST) {
-  StringRef Buf = AST.getSourceManager().getBufferData(
-      AST.getSourceManager().getMainFileID());
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
   // Edge-cases where the choice is forced.
   if (Buf.size() == 0)
     return {0, 0};
@@ -305,7 +272,6 @@ static std::pair<unsigned, unsigned> pointBounds(unsigned Offset,
 
 SelectionTree::SelectionTree(ASTContext &AST, unsigned Begin, unsigned End)
     : PrintPolicy(AST.getLangOpts()) {
-<<<<<<< HEAD
   // No fundamental reason the selection needs to be in the main file,
   // but that's all clangd has needed so far.
   FileID FID = AST.getSourceManager().getMainFileID();
@@ -320,16 +286,6 @@ SelectionTree::SelectionTree(ASTContext &AST, unsigned Begin, unsigned End)
 SelectionTree::SelectionTree(ASTContext &AST, unsigned Offset)
     : SelectionTree(AST, Offset, Offset) {}
 
-=======
-  if (Begin == End)
-    std::tie(Begin, End) = pointBounds(Begin, AST);
-  PrintPolicy.TerseOutput = true;
-
-  Nodes = SelectionVisitor(AST, Begin, End).take();
-  Root = Nodes.empty() ? nullptr : &Nodes.front();
-}
-
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
 const Node *SelectionTree::commonAncestor() const {
   if (!Root)
     return nullptr;
@@ -341,11 +297,5 @@ const Node *SelectionTree::commonAncestor() const {
   }
 }
 
-<<<<<<< HEAD
-=======
-SelectionTree::SelectionTree(ASTContext &AST, unsigned Offset)
-    : SelectionTree(AST, Offset, Offset) {}
-
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
 } // namespace clangd
 } // namespace clang
diff --git a/clangd/refactor/Tweak.cpp b/clangd/refactor/Tweak.cpp
index 02a22442..34634e64 100644
--- a/clangd/refactor/Tweak.cpp
+++ b/clangd/refactor/Tweak.cpp
@@ -38,6 +38,14 @@ void validateRegistry() {
 }
 } // namespace
 
+Tweak::Selection::Selection(ParsedAST &AST, unsigned RangeBegin,
+                            unsigned RangeEnd)
+    : AST(AST), ASTSelection(AST.getASTContext(), RangeBegin, RangeEnd) {
+  auto &SM = AST.getASTContext().getSourceManager();
+  Code = SM.getBufferData(SM.getMainFileID());
+  Cursor = SM.getComposedLoc(SM.getMainFileID(), RangeBegin);
+}
+
 std::vector<std::unique_ptr<Tweak>> prepareTweaks(const Tweak::Selection &S) {
   validateRegistry();
 
diff --git a/clangd/refactor/Tweak.h b/clangd/refactor/Tweak.h
index abd4995e..94164bc7 100644
--- a/clangd/refactor/Tweak.h
+++ b/clangd/refactor/Tweak.h
@@ -21,6 +21,7 @@
 
 #include "ClangdUnit.h"
 #include "Protocol.h"
+#include "Selection.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
@@ -39,16 +40,16 @@ class Tweak {
 public:
   /// Input to prepare and apply tweaks.
   struct Selection {
+    Selection(ParsedAST &AST, unsigned RangeBegin, unsigned RangeEnd);
     /// The text of the active document.
     llvm::StringRef Code;
     /// Parsed AST of the active file.
     ParsedAST &AST;
     /// A location of the cursor in the editor.
     SourceLocation Cursor;
-    // FIXME: add selection when there are checks relying on it.
+    // The AST nodes that were selected.
+    SelectionTree ASTSelection;
     // FIXME: provide a way to get sources and ASTs for other files.
-    // FIXME: cache some commonly required information (e.g. AST nodes under
-    //        cursor) to avoid redundant AST visit in every action.
   };
   virtual ~Tweak() = default;
   /// A unique id of the action, it is always equal to the name of the class
diff --git a/clangd/refactor/tweaks/SwapIfBranches.cpp b/clangd/refactor/tweaks/SwapIfBranches.cpp
index 2e0f33af..f28cbe53 100644
--- a/clangd/refactor/tweaks/SwapIfBranches.cpp
+++ b/clangd/refactor/tweaks/SwapIfBranches.cpp
@@ -41,58 +41,23 @@ public:
   std::string title() const override;
 
 private:
-  IfStmt *If = nullptr;
+  const IfStmt *If = nullptr;
 };
 
 REGISTER_TWEAK(SwapIfBranches);
 
-class FindIfUnderCursor : public RecursiveASTVisitor<FindIfUnderCursor> {
-public:
-  FindIfUnderCursor(ASTContext &Ctx, SourceLocation CursorLoc, IfStmt *&Result)
-      : Ctx(Ctx), CursorLoc(CursorLoc), Result(Result) {}
-
-  bool VisitIfStmt(IfStmt *If) {
-    // Check if the cursor is in the range of 'if (cond)'.
-    // FIXME: this does not contain the closing paren, add it too.
-    auto R = toHalfOpenFileRange(
-        Ctx.getSourceManager(), Ctx.getLangOpts(),
-        SourceRange(If->getIfLoc(), If->getCond()->getEndLoc().isValid()
-                                        ? If->getCond()->getEndLoc()
-                                        : If->getIfLoc()));
-    if (R && halfOpenRangeTouches(Ctx.getSourceManager(), *R, CursorLoc)) {
-      Result = If;
-      return false;
-    }
-    // Check the range of 'else'.
-    R = toHalfOpenFileRange(Ctx.getSourceManager(), Ctx.getLangOpts(),
-                            SourceRange(If->getElseLoc()));
-    if (R && halfOpenRangeTouches(Ctx.getSourceManager(), *R, CursorLoc)) {
-      Result = If;
+bool SwapIfBranches::prepare(const Selection &Inputs) {
+  for (const SelectionTree::Node *N = Inputs.ASTSelection.commonAncestor();
+       N && !If; N = N->Parent) {
+    // Stop once we hit a block, e.g. a lambda in the if condition.
+    if (dyn_cast_or_null<CompoundStmt>(N->ASTNode.get<Stmt>()))
       return false;
-    }
-
-    return true;
+    If = dyn_cast_or_null<IfStmt>(N->ASTNode.get<Stmt>());
   }
-
-private:
-  ASTContext &Ctx;
-  SourceLocation CursorLoc;
-  IfStmt *&Result;
-};
-} // namespace
-
-bool SwapIfBranches::prepare(const Selection &Inputs) {
-  auto &Ctx = Inputs.AST.getASTContext();
-  FindIfUnderCursor(Ctx, Inputs.Cursor, If).TraverseAST(Ctx);
-  if (!If)
-    return false;
-
   // avoid dealing with single-statement brances, they require careful handling
   // to avoid changing semantics of the code (i.e. dangling else).
-  if (!If->getThen() || !llvm::isa<CompoundStmt>(If->getThen()) ||
-      !If->getElse() || !llvm::isa<CompoundStmt>(If->getElse()))
-    return false;
-  return true;
+  return If && dyn_cast_or_null<CompoundStmt>(If->getThen()) &&
+         dyn_cast_or_null<CompoundStmt>(If->getElse());
 }
 
 Expected<tooling::Replacements> SwapIfBranches::apply(const Selection &Inputs) {
@@ -128,5 +93,7 @@ Expected<tooling::Replacements> SwapIfBranches::apply(const Selection &Inputs) {
 }
 
 std::string SwapIfBranches::title() const { return "Swap if branches"; }
+
+} // namespace
 } // namespace clangd
 } // namespace clang
diff --git a/unittests/clangd/TweakTests.cpp b/unittests/clangd/TweakTests.cpp
index 905f06b8..f747c991 100644
--- a/unittests/clangd/TweakTests.cpp
+++ b/unittests/clangd/TweakTests.cpp
@@ -52,9 +52,9 @@ void checkAvailable(StringRef ID, llvm::StringRef Input, bool Available) {
   ParsedAST AST = TU.build();
 
   auto CheckOver = [&](Range Selection) {
-    auto CursorLoc = llvm::cantFail(sourceLocationInMainFile(
-        AST.getASTContext().getSourceManager(), Selection.start));
-    auto T = prepareTweak(ID, Tweak::Selection{Code.code(), AST, CursorLoc});
+    unsigned Begin = cantFail(positionToOffset(Code.code(), Selection.start));
+    unsigned End = cantFail(positionToOffset(Code.code(), Selection.end));
+    auto T = prepareTweak(ID, Tweak::Selection(AST, Begin, End));
     if (Available)
       EXPECT_THAT_EXPECTED(T, Succeeded())
           << "code is " << markRange(Code.code(), Selection);
@@ -92,9 +92,9 @@ llvm::Expected<std::string> apply(StringRef ID, llvm::StringRef Input) {
   TU.Code = Code.code();
 
   ParsedAST AST = TU.build();
-  auto CursorLoc = llvm::cantFail(sourceLocationInMainFile(
-      AST.getASTContext().getSourceManager(), SelectionRng.start));
-  Tweak::Selection S = {Code.code(), AST, CursorLoc};
+  unsigned Begin = cantFail(positionToOffset(Code.code(), SelectionRng.start));
+  unsigned End = cantFail(positionToOffset(Code.code(), SelectionRng.end));
+  Tweak::Selection S(AST, Begin, End);
 
   auto T = prepareTweak(ID, S);
   if (!T)
@@ -149,6 +149,41 @@ TEST(TweakTest, SwapIfBranches) {
     }
   )cpp";
   checkTransform(ID, Input, Output);
+
+  // Available in subexpressions of the condition.
+  checkAvailable(ID, R"cpp(
+    void test() {
+      if(2 + [[2]] + 2) { return 2 + 2 + 2; } else { continue; }
+    }
+  )cpp");
+  // But not as part of the branches.
+  checkNotAvailable(ID, R"cpp(
+    void test() {
+      if(2 + 2 + 2) { return 2 + [[2]] + 2; } else { continue; }
+    }
+  )cpp");
+  // Range covers the "else" token, so available.
+  checkAvailable(ID, R"cpp(
+    void test() {
+      if(2 + 2 + 2) { return 2 + [[2 + 2; } else { continue;]] }
+    }
+  )cpp");
+  // Not available in compound statements in condition.
+  checkNotAvailable(ID, R"cpp(
+    void test() {
+      if([]{return [[true]];}()) { return 2 + 2 + 2; } else { continue; }
+    }
+  )cpp");
+  // Not available if both sides aren't braced.
+  checkNotAvailable(ID, R"cpp(
+    void test() {
+      ^if (1) return; else { return; }
+    }
+  )cpp");
+  // Only one if statement is supported!
+  checkNotAvailable(ID, R"cpp(
+    [[if(1){}else{}if(2){}else{}]]
+  )cpp");
 }
 
 } // namespace