RangeInfo: Treat comments transparently when resolving a given range. rdar://30926193 (#8052)

Author: nkcsgexi

lib/IDE/SwiftSourceDocInfo.cpp

@@ -19,6 +19,7 @@
 #include "swift/IDE/CommentConversion.h"
 #include "swift/IDE/Utils.h"
 #include "swift/Markup/XMLUtils.h"
+#include "swift/Subsystems.h"
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclObjC.h"
@@ -225,9 +226,11 @@ void ResolvedRangeInfo::print(llvm::raw_ostream &OS) {
     OS << "</Type>\n";
   }
 
-  OS << "<Context>";
-  printContext(OS, RangeContext);
-  OS << "</Context>\n";
+  if (RangeContext) {
+    OS << "<Context>";
+    printContext(OS, RangeContext);
+    OS << "</Context>\n";
+  }
 
   if (!HasSingleEntry) {
     OS << "<Entry>Multi</Entry>\n";
@@ -341,6 +344,9 @@ struct RangeResolver::Implementation {
       ContainedInRange(ContainedInRange) {}
   };
 
+  std::vector<Token> AllTokens;
+  Token &StartTok;
+  Token &EndTok;
   SourceLoc Start;
   SourceLoc End;
   StringRef Content;
@@ -412,43 +418,6 @@ struct RangeResolver::Implementation {
     return true;
   }
 
-  static SourceLoc getNonwhitespaceLocBefore(SourceManager &SM,
-                                             unsigned BufferID,
-                                             unsigned Offset) {
-    CharSourceRange entireRange = SM.getRangeForBuffer(BufferID);
-    StringRef Buffer = SM.extractText(entireRange);
-
-    const char *BufStart = Buffer.data();
-    if (Offset >= Buffer.size())
-      return SourceLoc();
-
-    for (unsigned Off = Offset; Off != 0; Off --) {
-      if (!clang::isWhitespace(*(BufStart + Off))) {
-        return SM.getLocForOffset(BufferID, Off);
-      }
-    }
-    return clang::isWhitespace(*BufStart) ? SourceLoc() :
-    SM.getLocForOffset(BufferID, 0);
-  }
-
-  static SourceLoc getNonwhitespaceLocAfter(SourceManager &SM,
-                                            unsigned BufferID,
-                                            unsigned Offset) {
-    CharSourceRange entireRange = SM.getRangeForBuffer(BufferID);
-    StringRef Buffer = SM.extractText(entireRange);
-
-    const char *BufStart = Buffer.data();
-    if (Offset >= Buffer.size())
-      return SourceLoc();
-
-    for (unsigned Off = Offset; Off < Buffer.size(); Off ++) {
-      if (!clang::isWhitespace(*(BufStart + Off))) {
-        return SM.getLocForOffset(BufferID, Off);
-      }
-    }
-    return SourceLoc();
-  }
-
   DeclContext *getImmediateContext() {
     for (auto It = ContextStack.rbegin(); It != ContextStack.rend(); It ++) {
       if (auto *DC = It->Parent.getAsDeclContext())
@@ -457,9 +426,14 @@ struct RangeResolver::Implementation {
     return static_cast<DeclContext*>(&File);
   }
 
-  Implementation(SourceFile &File, SourceLoc Start, SourceLoc End) :
-    File(File), Ctx(File.getASTContext()), SM(Ctx.SourceMgr), Start(Start),
-    End(End), Content(getContent()) {}
+  Implementation(SourceFile &File, std::vector<Token> AllTokens,
+                 unsigned StartIdx, unsigned EndIdx) :
+    File(File), Ctx(File.getASTContext()), SM(Ctx.SourceMgr),
+    AllTokens(AllTokens), StartTok(AllTokens[StartIdx]), EndTok(AllTokens[EndIdx]),
+    Start(StartTok.getLoc()), End(EndTok.getLoc()),
+    Content(getContent()) {
+      assert(Start.isValid() && End.isValid());
+  }
 
 public:
   bool hasResult() { return Result.hasValue(); }
@@ -490,14 +464,33 @@ struct RangeResolver::Implementation {
                                         unsigned Length) {
     SourceManager &SM = File.getASTContext().SourceMgr;
     unsigned BufferId = File.getBufferID().getValue();
-    SourceLoc StartLoc = Implementation::getNonwhitespaceLocAfter(SM, BufferId,
-                                                                  StartOff);
-    SourceLoc EndLoc = Implementation::getNonwhitespaceLocBefore(SM, BufferId,
-                                                         StartOff + Length - 1);
-    StartLoc = Lexer::getLocForStartOfToken(SM, StartLoc);
-    EndLoc = Lexer::getLocForStartOfToken(SM, EndLoc);
-    return StartLoc.isInvalid() || EndLoc.isInvalid() ? nullptr :
-      new Implementation(File, StartLoc, EndLoc);
+
+    LangOptions Opts = File.getASTContext().LangOpts;
+    Opts.AttachCommentsToDecls = true;
+    std::vector<Token> AllTokens = tokenize(Opts, SM, BufferId, 0, 0, false);
+    auto TokenComp = [&](Token &LHS, SourceLoc Loc) {
+      return SM.isBeforeInBuffer(LHS.getLoc(), Loc);
+    };
+
+    SourceLoc StartRaw = SM.getLocForOffset(BufferId, StartOff);
+    SourceLoc EndRaw = SM.getLocForOffset(BufferId, StartOff + Length);
+
+    // This points to the first token after or on the start loc.
+    auto StartIt = std::lower_bound(AllTokens.begin(), AllTokens.end(), StartRaw,
+                                    TokenComp);
+    // This points to the first token after or on the end loc;
+    auto EndIt = std::lower_bound(AllTokens.begin(), AllTokens.end(), EndRaw,
+                                  TokenComp);
+    // Erroneous case.
+    if (StartIt == AllTokens.end() || EndIt == AllTokens.begin())
+      return nullptr;
+
+    // The start token is inclusive.
+    unsigned StartIdx = StartIt - AllTokens.begin();
+
+    // The end token is exclusive.
+    unsigned EndIdx = EndIt - 1 - AllTokens.begin();
+    return new Implementation(File, std::move(AllTokens), StartIdx, EndIdx);
   }
 
   static Implementation *createInstance(SourceFile &File, SourceLoc Start,
@@ -728,7 +721,9 @@ struct RangeResolver::Implementation {
 
   StringRef getContent() {
     SourceManager &SM = File.getASTContext().SourceMgr;
-    return CharSourceRange(SM, Start, Lexer::getLocForEndOfToken(SM, End)).str();
+    return CharSourceRange(SM, StartTok.hasComment() ?
+                            StartTok.getCommentStart() : StartTok.getLoc(),
+                           Lexer::getLocForEndOfToken(SM, End)).str();
   }
 };
 

test/IDE/range_info_comments.swift

@@ -0,0 +1,46 @@
+func foo() -> Int{
+  // some comments
+  var aaa = 1 + 2
+  aaa = aaa + 3
+  if aaa == 3 { aaa = 4 }
+  // some comments
+  return aaa
+}
+
+func foo1() -> Int{
+  /// some comments
+  var aaa = 1 + 2
+  aaa = aaa + 3
+  if aaa == 3 { aaa = 4 }
+  /// some comments
+  return aaa
+}
+
+func foo2() -> Int{
+  /* some comments*/
+  var aaa = 1 + 2
+  aaa = aaa + 3
+  if aaa == 3 { aaa = 4 }
+  /* some comments*/
+  return aaa
+}
+
+// RUN: %target-swift-ide-test -range -pos=2:1 -end-pos 6:19 -source-filename %s | %FileCheck %s -check-prefix=CHECK1
+// RUN: %target-swift-ide-test -range -pos=11:1 -end-pos 15:20 -source-filename %s | %FileCheck %s -check-prefix=CHECK-KIND
+// RUN: %target-swift-ide-test -range -pos=20:1 -end-pos 24:21 -source-filename %s | %FileCheck %s -check-prefix=CHECK-KIND
+// RUN: %target-swift-ide-test -range -pos=1:1 -end-pos 15:20 -source-filename %s | %FileCheck %s -check-prefix=CHECK-INVALID
+
+// CHECK1: <Kind>MultiStatement</Kind>
+// CHECK1-NEXT: <Content>// some comments
+// CHECK1-NEXT:   var aaa = 1 + 2
+// CHECK1-NEXT:   aaa = aaa + 3
+// CHECK1-NEXT:   if aaa == 3 { aaa = 4 }</Content>
+// CHECK1-NEXT: <Type>Void</Type>
+// CHECK1-NEXT: <Context>swift_ide_test.(file).foo()</Context>
+// CHECK1-NEXT: <Declared>aaa</Declared><OutscopeReference>true</OutscopeReference>
+// CHECK1-NEXT: <Referenced>aaa</Referenced><Type>@lvalue Int</Type>
+// CHECK1-NEXT: <ASTNodes>3</ASTNodes>
+// CHECK1-NEXT: <end>
+
+// CHECK-KIND: <Kind>MultiStatement</Kind>
+// CHECK-INVALID: <Kind>Invalid</Kind>

test/SourceKit/RangeInfo/basic.swift

@@ -17,12 +17,9 @@ struct S { func foo() {} }
 // RUN: %sourcekitd-test -req=range -pos=4:1 -length 25 %s -- %s | %FileCheck %s -check-prefix=CHECK3
 // RUN: %sourcekitd-test -req=range -pos=4:1 -length 26 %s -- %s | %FileCheck %s -check-prefix=CHECK3
 // RUN: %sourcekitd-test -req=range -pos=4:1 -length 27 %s -- %s | %FileCheck %s -check-prefix=CHECK3
-// RUN: %sourcekitd-test -req=range -pos=4:4 -length 22 %s -- %s | %FileCheck %s -check-prefix=CHECK3
 
 // RUN: %sourcekitd-test -req=range -pos=5:1 -length 12 %s -- %s | %FileCheck %s -check-prefix=CHECK4
 // RUN: %sourcekitd-test -req=range -pos=5:2 -length 11 %s -- %s | %FileCheck %s -check-prefix=CHECK4
-// RUN: %sourcekitd-test -req=range -pos=5:5 -length 8 %s -- %s | %FileCheck %s -check-prefix=CHECK4
-// RUN: %sourcekitd-test -req=range -pos=5:5 -length 9 %s -- %s | %FileCheck %s -check-prefix=CHECK4
 
 // RUN: %sourcekitd-test -req=range -pos=8:1 -length 31 %s -- %s | %FileCheck %s -check-prefix=CHECK5
 // RUN: %sourcekitd-test -req=range -pos=9:1 -length 25 %s -- %s | %FileCheck %s -check-prefix=CHECK6