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

[sammccall]

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).

[sammccall]

No real urgency here (though it may simplify the pending tweaks) - mostly this is motivation for D57562.

Preserving the ability to fire on the condition and its subexpressions, but not the if/else branches, is pretty arbitrary and mostly serves to demontrate a nontrivial traversal.

The selection-matching code is shorter but more importantly I think it's a bit less error-prone - it should be harder to introduce off-by-one errors and index miscalculations without SourceLocations everywhere.

[ilya-biryukov]

I wonder if this makes the check available in too many places?

void test() {
  if (callFunc([]() { 
         for (int i = 0; i < 100; ++i) { /* compute something */; } // action also available here...
      }) {
   } else {} 
}

This look a bit artificial for this check as conditions are rarely large, but could be more significant for other checks, where it could be more common to have declarations of classes, functions and lambdas lambdas recursively inside other declarations.
Maybe we should consider stopping the traversals up at some points, e.g. at lambda decls or something similar to limit this class of problems in the future?

clangd/refactor/Tweak.cpp

43 ↗

(On Diff #184690)

NIT: maybe move the code into the ctor body and save AST.getASTContext() and getSourceManager() into local vars for readability?
There's too much noise with those getters currently.

unittests/clangd/TweakTests.cpp

56 ↗

(On Diff #184690)

NIT: split into two declarations?

96 ↗

(On Diff #184690)

NIT: split into two declaration?

[sammccall]

In D57570#1380216, @ilya-biryukov wrote:

I wonder if this makes the check available in too many places?

void test() {
  if (callFunc([]() { 
         for (int i = 0; i < 100; ++i) { /* compute something */; } // action also available here...
      }) {
   } else {} 
}

This look a bit artificial for this check as conditions are rarely large, but could be more significant for other checks, where it could be more common to have declarations of classes, functions and lambdas lambdas recursively inside other declarations.
Maybe we should consider stopping the traversals up at some points, e.g. at lambda decls or something similar to limit this class of problems in the future?

Heh, this actually makes it easier to implement - we can break at a compoundstmt, and since we require both if and then clauses to be braced, this takes care of only alllowing subexpressions of cond.

Maybe we should consider stopping the traversals up at some points, e.g. at lambda decls or something similar to limit this class of problems in the future?

If you mean in some frameworky sort of way... yeah, it's quite likely we need something. At least to make the Node* -> DynTypedNode -> Stmt -> IfStmt dance less clunky.
I've deliberately avoided adding too much in the way of helpers for now in favor of the raw data structures, mostly wanting to see what the common patterns are - we can afford to write a few of these by hand.

clangd/refactor/Tweak.cpp

43 ↗

(On Diff #184690)

Whoops, much better.
I forgot *some* of these can be default-constructed.

unittests/clangd/TweakTests.cpp

56 ↗

(On Diff #184690)

(clearly we should chat about the two-decls thing :-) on the other review?)

[sammccall]

Simplify Selection constructor.
Don't traverse up through compoundstmt.
Add more tests.

[ilya-biryukov]

In D57570#1380327, @sammccall wrote:

Maybe we should consider stopping the traversals up at some points, e.g. at lambda decls or something similar to limit this class of problems in the future?

If you mean in some frameworky sort of way... yeah, it's quite likely we need something. At least to make the Node* -> DynTypedNode -> Stmt -> IfStmt dance less clunky.
I've deliberately avoided adding too much in the way of helpers for now in favor of the raw data structures, mostly wanting to see what the common patterns are - we can afford to write a few of these by hand.

Exactly, the code dealing with typed nodes there is definitely too verbose. Thinking about helpers like findAncestor<IfStmt>().
And yeah, totally fine to keep it as is for now, it's great that the change is focused.

Another concern that came to mind: what should we do if two if statements are selected? IIUC the commonAncestor would be their parent and we don't have the action available, right? This LG, the alternatives are showing multiple tweaks or just one of them and inevitably confusing the users in some cases (either applying the tweak to the other than expected if statement or showing two tweaks with identical descriptions and no ways to distinguish between those).
Could we add a test for this case?

unittests/clangd/TweakTests.cpp

56 ↗

(On Diff #184690)

Done :-)

[ilya-biryukov]

clangd/refactor/tweaks/SwapIfBranches.cpp

60 ↗

(On Diff #184707)

NIT: drop llvm:: for consistency.

[ilya-biryukov]

LGTM.
See the NIT about llvm::, though. Also adding a test with multiple selected ifs would be reasonable, to document and validate behaviour in such cases.

[sammccall]

Split declarations, add test for two-if-statements.

[sammccall]

In D57570#1380354, @ilya-biryukov wrote:

In D57570#1380327, @sammccall wrote:

Maybe we should consider stopping the traversals up at some points, e.g. at lambda decls or something similar to limit this class of problems in the future?

If you mean in some frameworky sort of way... yeah, it's quite likely we need something. At least to make the Node* -> DynTypedNode -> Stmt -> IfStmt dance less clunky.
I've deliberately avoided adding too much in the way of helpers for now in favor of the raw data structures, mostly wanting to see what the common patterns are - we can afford to write a few of these by hand.

Exactly, the code dealing with typed nodes there is definitely too verbose. Thinking about helpers like findAncestor<IfStmt>().
And yeah, totally fine to keep it as is for now, it's great that the change is focused.

e.g. findAncestor is nice but here we want to walk up but break if we hit a limit node (e.g. compoundstmt).
If that turns out to be the normal case, then a plain findAncestor isn't such a useful primitive.
Some other thoughts:

• In a draft I did have a tree traversal where the node visitor could return {continue, abort, prune} - seemed cool but I didn't actually have a use for it
• just some direct casting operators (Node->as<IfStmt>()) would help, probably.

Another concern that came to mind: what should we do if two if statements are selected? IIUC the commonAncestor would be their parent and we don't have the action available, right?

Correct. Ranges are confusing, I think we need to mostly treat them as "a target to perform an action on" rather than "a selection of targets to perform actions on". Maybe some exceptions like maybe bulk actions on includes or member function decls, but the simple model should work for most tweaks.

Could we add a test for this case?

Done.

clang-tools-extra/trunk/clangd/ClangdServer.cpp

auto Action = [Pos](Path File, std::string NewName,
     }
     return CB(std::move(Replacements));
   };
 
   WorkScheduler.runWithAST(
       "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));
   };
 
   WorkScheduler.runWithAST("EnumerateTweaks", File,
                            Bind(Action, std::move(CB), File.str()));
 }
 
 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
                               Callback<tooling::Replacements> CB) {
   auto Action = [Sel](decltype(CB) CB, std::string File, std::string 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,
       Bind(Action, std::move(CB), File.str(), TweakID.str()));
 }
 
 void ClangdServer::dumpAST(PathRef File,
                            llvm::unique_function<void(std::string)> Callback) {

clang-tools-extra/trunk/clangd/Selection.cpp

bool traverseNode(T *Node, const Func &Body) {
     bool Ret = Body();
     pop();
     return Ret;
   }
 
   // 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;
   }
 
   // Pushes a node onto the ancestor stack. Pairs with pop().
   void push(DynTypedNode Node) {
     Nodes.emplace_back();
     Nodes.back().ASTNode = std::move(Node);
SelectionTree::Selection computeSelection(const Node &N) {
     if (!S.isValid())
       return SelectionTree::Unselected;
     // getTopMacroCallerLoc() allows selection of constructs in macro args. e.g:
     //   #define LOOP_FOREVER(Body) for(;;) { Body }
     //   void IncrementLots(int &x) {
     //     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.
     if (B.first != SelFile || E.first != SelFile)
       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" 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;
 
     // We hit something, need some more precise checks.
     // Adjust [B, E) to be a half-open character range.
     E.second += Lexer::MeasureTokenLength(S.getEnd(), SM, LangOpts);
     // This node's own selected text is (this range ^ selection) - child ranges.
     // If that's empty, then we've only collided with children.
if (Nodes.empty())
       return false;
 
     // Collect all the expansion ranges, as offsets.
     SmallVector<std::pair<unsigned, unsigned>, 8> ChildRanges;
     for (const Node *N : Nodes) {
       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.
       if (auto Tok = Lexer::findNextToken(R.getEnd(), SM, LangOpts))
         E.second = SM.getFileOffset(Tok->getLocation());
       else
         E.second += Lexer::MeasureTokenLength(R.getEnd(), SM, LangOpts);
       ChildRanges.push_back({B.second, E.second});
for (const auto R : ChildRanges) {
         return true; // Remaining range is empty.
     }
     return false; // Went through all children, trailing characters remain.
   }
 
   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;
   unsigned SelEnd;
   FileID SelFile;
   // If the selection start slices a token in half, the beginning of that token.
   // This is useful for checking whether the end of a token range overlaps
   // the selection: range.end < SelBeginTokenStart is equivalent to
   // 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
 
 void SelectionTree::print(llvm::raw_ostream &OS, const SelectionTree::Node &N,
                           int Indent) const {
   if (N.Selected)
     OS.indent(Indent - 1) << (N.Selected == SelectionTree::Complete ? '*'
                                                                     : '.');
   else
     OS.indent(Indent);
   OS << N.ASTNode.getNodeKind().asStringRef() << " ";
   N.ASTNode.print(OS, PrintPolicy);
   OS << "\n";
   for (const Node *Child : N.Children)
     print(OS, *Child, Indent + 2);
 }
 
 // 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};
   if (Offset == 0)
     return {0, 1};
   if (Offset == Buf.size())
     return {Offset - 1, Offset};
   // We could choose either this byte or the previous. Usually we prefer the
   // character on the right of the cursor (or under a block cursor).
   // But if that's whitespace, we likely want the token on the left.
   if (isWhitespace(Buf[Offset]) && !isWhitespace(Buf[Offset - 1]))
     return {Offset - 1, Offset};
   return {Offset, Offset + 1};
 }
 
 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();
   if (Begin == End)
     std::tie(Begin, End) = pointBounds(Begin, FID, AST);
   PrintPolicy.TerseOutput = true;
 
   Nodes = SelectionVisitor::collect(AST, Begin, End, FID);
   Root = Nodes.empty() ? nullptr : &Nodes.front();
 }
 
 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;
   for (const Node *Ancestor = Root;; Ancestor = Ancestor->Children.front()) {
     if (Ancestor->Selected || Ancestor->Children.size() > 1)
       return Ancestor;
     // The tree only contains ancestors of the interesting nodes.
     assert(!Ancestor->Children.empty() && "bad node in selection tree");
   }
 }
 
-<<<<<<< HEAD
-=======
-SelectionTree::SelectionTree(ASTContext &AST, unsigned Offset)
-    : SelectionTree(AST, Offset, Offset) {}
-
->>>>>>> [clangd] Lib to compute and represent selection under cursor.
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/clangd/refactor/Tweak.h

 //     It is performed when the user actually chooses the action.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_REFACTOR_ACTIONS_TWEAK_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_REFACTOR_ACTIONS_TWEAK_H
 
 #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"
 namespace clang {
 namespace clangd {
 
 /// An interface base for small context-sensitive refactoring actions.
 /// To implement a new tweak use the following pattern in a .cpp file:
 ///   class MyTweak : public Tweak {
 ///   public:
 ///     const char* id() const override final; // defined by REGISTER_TWEAK.
 ///     // implement other methods here.
 ///   };
 ///   REGISTER_TWEAK(MyTweak);
 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 *
     /// 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
   /// defining the Tweak. Definition is provided automatically by
   /// REGISTER_TWEAK.
   virtual const char *id() const = 0;
   /// Run the first stage of the action. The non-None result indicates that the
   /// action is available and should be shown to the user. Returns None if the

clang-tools-extra/trunk/clangd/refactor/Tweak.cpp

for (const auto &E : TweakRegistry::entries()) {
     assert(E.instantiate()->id() == E.getName() &&
            "id should be equal to class name");
     assert(Seen.try_emplace(E.getName()).second && "duplicate check id");
   }
 #endif
 }
 } // 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();
 
   std::vector<std::unique_ptr<Tweak>> Available;
   for (const auto &E : TweakRegistry::entries()) {
     std::unique_ptr<Tweak> T = E.instantiate();
     if (!T->prepare(S))
       continue;

clang-tools-extra/trunk/clangd/refactor/tweaks/SwapIfBranches.cpp

 public:
   const char *id() const override final;
 
   bool prepare(const Selection &Inputs) override;
   Expected<tooling::Replacements> apply(const Selection &Inputs) override;
   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;
-      return false;
-    }
-
-    return true;
-  }
-
-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)
+  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;
-
+    If = dyn_cast_or_null<IfStmt>(N->ASTNode.get<Stmt>());
+  }
   // 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) {
   auto &Ctx = Inputs.AST.getASTContext();
   auto &SrcMgr = Ctx.getSourceManager();
 
   auto ThenRng = toHalfOpenFileRange(SrcMgr, Ctx.getLangOpts(),
                                      If->getThen()->getSourceRange());
Expected<tooling::Replacements> SwapIfBranches::apply(const Selection &Inputs) {
   if (auto Err = Result.add(tooling::Replacement(Ctx.getSourceManager(),
                                                  ElseRng->getBegin(),
                                                  ElseCode.size(), ThenCode)))
     return std::move(Err);
   return Result;
 }
 
 std::string SwapIfBranches::title() const { return "Swap if branches"; }
+
+} // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/unittests/clangd/TweakTests.cpp

ASSERT_TRUE(0 < Code.points().size() || 0 < Code.ranges().size())
       << "no points of interest specified";
   TestTU TU;
   TU.Filename = "foo.cpp";
   TU.Code = Code.code();
 
   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);
     else
       EXPECT_THAT_EXPECTED(T, Failed())
           << "code is " << markRange(Code.code(), Selection);
   };
   for (auto P : Code.points())
llvm::Expected<std::string> apply(StringRef ID, llvm::StringRef Input) {
   } else {
     SelectionRng = Code.range();
   }
   TestTU TU;
   TU.Filename = "foo.cpp";
   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)
     return T.takeError();
   auto Replacements = (*T)->apply(S);
   if (!Replacements)
     return Replacements.takeError();
   return applyAllReplacements(Code.code(), *Replacements);
Input = R"cpp(
     }
   )cpp";
   Output = R"cpp(
     void test() {
       if () { continue; } else { return 100; }
     }
   )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
 } // namespace clangd
 } // namespace clang