[SourceKit] Consider source location in the orignal buffer when computing refactoring actions

Author: ahoppen

include/swift/IDE/Utils.h

@@ -139,6 +139,9 @@ struct ResolvedCursorInfo {
   CursorInfoKind Kind = CursorInfoKind::Invalid;
   SourceFile *SF = nullptr;
   SourceLoc Loc;
+  /// If \c Loc resides in a temporary buffer during AST reuse, the
+  /// corresponding location in the original, full buffer.
+  SourceLoc LocInOriginalBuffer;
 
   // Technically, these structs could form a union (because only one of them is
   // active at a time). But I had issues with C++ complaining about copy
@@ -185,7 +188,11 @@ struct ResolvedCursorInfo {
   SourceFile *getSourceFile() const { return SF; }
 
   SourceLoc getLoc() const { return Loc; }
-  void setLoc(SourceLoc Loc) { this->Loc = Loc; }
+  SourceLoc getLocInOriginalBuffer() const { return LocInOriginalBuffer; }
+  void setLoc(SourceLoc Loc, SourceLoc LocInOriginalBuffer) {
+    this->Loc = Loc;
+    this->LocInOriginalBuffer = LocInOriginalBuffer;
+  }
 
   friend bool operator==(const ResolvedCursorInfo &lhs,
                          const ResolvedCursorInfo &rhs) {

lib/IDE/CursorInfo.cpp

@@ -247,7 +247,7 @@ class CursorInfoDoneParsingCallback : public IDEInspectionCallbacks {
         /*CtorTyRef=*/nullptr,
         /*ExtTyRef=*/nullptr, /*IsRef=*/false, /*Ty=*/Type(),
         /*ContainerType=*/Type());
-    CursorInfo->setLoc(RequestedLoc);
+    CursorInfo->setLoc(RequestedLoc, RequestedLocInOriginalBuffer);
     CursorInfo->setShorthandShadowedDecls(
         Finder.getShorthandShadowedDecls(DeclResult->getDecl()));
     return CursorInfo;

lib/IDE/IDERequests.cpp

@@ -175,7 +175,7 @@ bool CursorInfoResolver::visitSubscriptReference(ValueDecl *D,
 ResolvedCursorInfo CursorInfoResolver::resolve(SourceLoc Loc) {
   assert(Loc.isValid());
   LocToResolve = Loc;
-  CursorInfo.setLoc(Loc);
+  CursorInfo.setLoc(Loc, /*LocInOriginalBuffer=*/Loc);
 
   walk(SrcFile);
 

lib/Refactoring/Refactoring.cpp

@@ -49,27 +49,33 @@ class ContextFinder : public SourceEntityWalker {
   ASTContext &Ctx;
   SourceManager &SM;
   SourceRange Target;
-  std::function<bool(ASTNode)> IsContext;
+  SourceRange TargetInOriginalBuffer;
   SmallVector<ASTNode, 4> AllContexts;
   bool contains(ASTNode Enclosing) {
-    auto Result = SM.rangeContains(Enclosing.getSourceRange(), Target);
-    if (Result && IsContext(Enclosing))
+    auto Result =
+        SM.rangeContains(Enclosing.getSourceRange(), Target) ||
+        SM.rangeContains(Enclosing.getSourceRange(), TargetInOriginalBuffer);
+    if (Result && IsContext(Enclosing)) {
       AllContexts.push_back(Enclosing);
+    }
     return Result;
   }
 public:
-  ContextFinder(SourceFile &SF, ASTNode TargetNode,
-                std::function<bool(ASTNode)> IsContext =
-                  [](ASTNode N) { return true; }) :
-                  SF(SF), Ctx(SF.getASTContext()), SM(Ctx.SourceMgr),
-                  Target(TargetNode.getSourceRange()), IsContext(IsContext) {}
-  ContextFinder(SourceFile &SF, SourceLoc TargetLoc,
-                std::function<bool(ASTNode)> IsContext =
-                  [](ASTNode N) { return true; }) :
-                  SF(SF), Ctx(SF.getASTContext()), SM(Ctx.SourceMgr),
-                  Target(TargetLoc), IsContext(IsContext) {
-                    assert(TargetLoc.isValid() && "Invalid loc to find");
-                  }
+  ContextFinder(
+      SourceFile &SF, ASTNode TargetNode,
+      std::function<bool(ASTNode)> IsContext = [](ASTNode N) { return true; })
+      : SF(SF), Ctx(SF.getASTContext()), SM(Ctx.SourceMgr),
+        Target(TargetNode.getSourceRange()),
+        TargetInOriginalBuffer(TargetNode.getSourceRange()),
+        IsContext(IsContext) {}
+  ContextFinder(
+      SourceFile &SF, SourceLoc TargetLoc, SourceLoc TargetLocInOriginalBuffer,
+      std::function<bool(ASTNode)> IsContext = [](ASTNode N) { return true; })
+      : SF(SF), Ctx(SF.getASTContext()), SM(Ctx.SourceMgr), Target(TargetLoc),
+        TargetInOriginalBuffer(TargetLocInOriginalBuffer),
+        IsContext(IsContext) {
+    assert(TargetLoc.isValid() && "Invalid loc to find");
+  }
   bool walkToDeclPre(Decl *D, CharSourceRange Range) override { return contains(D); }
   bool walkToStmtPre(Stmt *S) override { return contains(S); }
   bool walkToExprPre(Expr *E) override { return contains(E); }
@@ -3928,10 +3934,11 @@ static CallExpr *findTrailingClosureTarget(
     return nullptr;
 
   // Find inner most CallExpr
-  ContextFinder Finder(
-      *CursorInfo.getSourceFile(), CursorInfo.getLoc(), [](ASTNode N) {
-        return N.isStmt(StmtKind::Brace) || N.isExpr(ExprKind::Call);
-      });
+  ContextFinder Finder(*CursorInfo.getSourceFile(), CursorInfo.getLoc(),
+                       CursorInfo.getLocInOriginalBuffer(), [](ASTNode N) {
+                         return N.isStmt(StmtKind::Brace) ||
+                                N.isExpr(ExprKind::Call);
+                       });
   Finder.resolve();
   auto contexts = Finder.getContexts();
   if (contexts.empty())
@@ -4165,7 +4172,7 @@ CallExpr *findOuterCall(const ResolvedCursorInfo &CursorInfo) {
   //       already generated a slice of the AST for anything that contains
   //       the cursor location
   ContextFinder Finder(*CursorInfo.getSourceFile(), CursorInfo.getLoc(),
-                       IncludeInContext);
+                       CursorInfo.getLocInOriginalBuffer(), IncludeInContext);
   Finder.resolve();
   auto Contexts = Finder.getContexts();
   if (Contexts.empty())
@@ -4191,7 +4198,7 @@ FuncDecl *findFunction(const ResolvedCursorInfo &CursorInfo) {
   };
 
   ContextFinder Finder(*CursorInfo.getSourceFile(), CursorInfo.getLoc(),
-                       IncludeInContext);
+                       CursorInfo.getLocInOriginalBuffer(), IncludeInContext);
   Finder.resolve();
 
   auto Contexts = Finder.getContexts();
@@ -8282,9 +8289,10 @@ bool RefactoringActionConvertCallToAsyncAlternative::performChange() {
          "Should not run performChange when refactoring is not applicable");
 
   // Find the scope this call is in
-  ContextFinder Finder(
-      *CursorInfo.getSourceFile(), CursorInfo.getLoc(),
-      [](ASTNode N) { return N.isStmt(StmtKind::Brace) && !N.isImplicit(); });
+  ContextFinder Finder(*CursorInfo.getSourceFile(), CursorInfo.getLoc(),
+                       CursorInfo.getLocInOriginalBuffer(), [](ASTNode N) {
+                         return N.isStmt(StmtKind::Brace) && !N.isImplicit();
+                       });
   Finder.resolve();
   auto Scopes = Finder.getContexts();
   BraceStmt *Scope = nullptr;
@@ -8437,7 +8445,8 @@ findMacroExpansionTargetExpr(const ResolvedCursorInfo &Info) {
   if (auto refInfo = dyn_cast<ResolvedValueRefCursorInfo>(&Info)) {
     if (refInfo->isRef() && isa_and_nonnull<MacroDecl>(refInfo->getValueD())) {
       ContextFinder Finder(
-          *Info.getSourceFile(), Info.getLoc(), [&](ASTNode N) {
+          *Info.getSourceFile(), Info.getLoc(), Info.getLocInOriginalBuffer(),
+          [&](ASTNode N) {
             auto *expr =
                 dyn_cast_or_null<MacroExpansionExpr>(N.dyn_cast<Expr *>());
             return expr &&

test/SourceKit/CursorInfo/refactoring_actions_with_ast_reuse.swift

@@ -0,0 +1,20 @@
+func materialize() {
+ let a = 2
+}
+
+func outer() {
+  func inner() {
+  }
+}
+
+// RUN: %sourcekitd-test -req=cursor -pos=2:6 %s -- %s == \
+// RUN: -req=cursor -req-opts=retrieve_refactor_actions=1 -pos=6:8 %s -- %s | %FileCheck %s
+
+// CHECK: source.lang.swift.decl.function.free (6:8-6:15)
+// CHECK: inner()
+// CHECK: ACTIONS BEGIN
+// CHECK-DAG: source.refactoring.kind.rename.local
+// CHECK-DAG: Local Rename
+// CHECK-DAG: source.refactoring.kind.convert.func-to-async
+// CHECK-DAG: Convert Function to Async
+// CHECK: ACTIONS END