Refactoring tool: Avoid re-tokenizing when collecting the available refactorings in a given source range. NFC (#11857)

Author: nkcsgexi

include/swift/Parse/Lexer.h

@@ -536,6 +536,16 @@ class Lexer {
   bool tryLexConflictMarker();
 };
 
+/// Given an ordered token \param Array , get the iterator pointing to the first
+/// token that is not before \param Loc .
+template<typename ArrayTy, typename Iterator = typename ArrayTy::iterator>
+Iterator token_lower_bound(ArrayTy &Array, SourceLoc Loc) {
+  return std::lower_bound(Array.begin(), Array.end(), Loc,
+    [](const Token &T, SourceLoc L) {
+      return T.getLoc().getOpaquePointerValue() < L.getOpaquePointerValue();
+  });
+}
+
 } // end namespace swift
 
 #endif

lib/IDE/Refactoring.cpp

@@ -1257,15 +1257,10 @@ struct SimilarExprCollector: public SourceEntityWalker {
 
   /// Find all tokens included by an expression.
   llvm::ArrayRef<Token> getExprSlice(Expr *E) {
-    auto TokenComp = [&](const Token &LHS, SourceLoc Loc) {
-      return SM.isBeforeInBuffer(LHS.getLoc(), Loc);
-    };
     SourceLoc StartLoc = E->getStartLoc();
     SourceLoc EndLoc = E->getEndLoc();
-    auto StartIt = std::lower_bound(AllTokens.begin(), AllTokens.end(),
-                                    StartLoc, TokenComp);
-    auto EndIt = std::lower_bound(AllTokens.begin(), AllTokens.end(),
-                                  EndLoc, TokenComp);
+    auto StartIt = token_lower_bound(AllTokens, StartLoc);
+    auto EndIt = token_lower_bound(AllTokens, EndLoc);
     assert(StartIt->getLoc() == StartLoc);
     assert(EndIt->getLoc() == EndLoc);
     return AllTokens.slice(StartIt - AllTokens.begin(), EndIt - StartIt + 1);
@@ -1333,16 +1328,16 @@ bool RefactoringActionExtractExprBase::performChange() {
     Collector.walk(BS);
 
     if (ExtractRepeated) {
-      unsigned BufferId = *TheFile->getBufferID();
-
       // Tokenize the brace statement; all expressions should have their tokens
       // in this array.
-      std::vector<Token> AllToks(tokenize(Ctx.LangOpts, SM, BufferId,
-          /*start offset*/SM.getLocOffsetInBuffer(BS->getStartLoc(), BufferId),
-          /*end offset*/SM.getLocOffsetInBuffer(BS->getEndLoc(), BufferId)));
+      auto AllTokens = TheFile->getAllTokens();
+      auto StartIt = token_lower_bound(AllTokens, BS->getStartLoc());
+      auto EndIt = token_lower_bound(AllTokens, BS->getEndLoc());
 
       // Collect all expressions we are going to extract.
-      SimilarExprCollector(SM, SelectedExpr, AllToks, AllExpressions).walk(BS);
+      SimilarExprCollector(SM, SelectedExpr,
+        AllTokens.slice(StartIt - AllTokens.begin(), EndIt - StartIt + 1),
+                           AllExpressions).walk(BS);
     } else {
       AllExpressions.insert(SelectedExpr);
     }

lib/IDE/SwiftSourceDocInfo.cpp

@@ -875,7 +875,7 @@ struct RangeResolver::Implementation {
   };
 
 
-  std::vector<Token> TokensInRange;
+  ArrayRef<Token> TokensInRange;
   SourceLoc Start;
   SourceLoc End;
 
@@ -1066,32 +1066,41 @@ struct RangeResolver::Implementation {
   createInstance(SourceFile &File, unsigned StartOff, unsigned Length) {
     SourceManager &SM = File.getASTContext().SourceMgr;
     unsigned BufferId = File.getBufferID().getValue();
-
-    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);
-    };
-
+    auto AllTokens = File.getAllTokens();
     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);
+    auto StartIt = token_lower_bound(AllTokens, StartRaw);
+
+    // Skip all the comments.
+    while(StartIt != AllTokens.end()) {
+      if (StartIt->getKind() != tok::comment)
+        break;
+      StartIt ++;
+    }
+
+    // Erroneous case.
+    if (StartIt == AllTokens.end())
+      return nullptr;
+
     // This points to the first token after or on the end loc;
-    auto EndIt = std::lower_bound(AllTokens.begin(), AllTokens.end(), EndRaw,
-                                  TokenComp);
+    auto EndIt = token_lower_bound(AllTokens, EndRaw);
+
+    // Adjust end token to skip comments.
+    while (EndIt != AllTokens.begin()) {
+      EndIt --;
+      if (EndIt->getKind() != tok::comment)
+        break;
+    }
+
     // Erroneous case.
-    if (StartIt == AllTokens.end() || EndIt == AllTokens.begin())
+    if (EndIt < StartIt)
       return nullptr;
 
-    // The start token is inclusive.
     unsigned StartIdx = StartIt - AllTokens.begin();
-
     return std::unique_ptr<Implementation>(new Implementation(File,
-      llvm::makeArrayRef(AllTokens.data() + StartIdx, EndIt - StartIt)));
+      AllTokens.slice(StartIdx, EndIt - StartIt + 1)));
   }
 
   static std::unique_ptr<Implementation>

lib/Parse/Parser.cpp

@@ -308,10 +308,7 @@ class TokenRecorder: public ConsumeTokenReceiver {
   llvm::DenseMap<const void*, tok> TokenKindChangeMap;
 
   std::vector<Token>::iterator lower_bound(SourceLoc Loc) {
-    return std::lower_bound(Bag.begin(), Bag.end(), Loc,
-      [](const Token &T, SourceLoc L) {
-        return T.getLoc().getOpaquePointerValue() < L.getOpaquePointerValue();
-    });
+    return token_lower_bound(Bag, Loc);
   }
 
   std::vector<Token>::iterator lower_bound(Token Tok) {
@@ -400,7 +397,6 @@ class TokenRecorder: public ConsumeTokenReceiver {
     llvm::SmallVector<Token, 4> TokensToConsume;
     if (Tok.hasComment()) {
       relexComment(Tok.getCommentRange(), TokensToConsume);
-      Tok.clearCommentLength();
     }
 
     TokensToConsume.push_back(Tok);

tools/swift-ide-test/swift-ide-test.cpp

@@ -2669,6 +2669,7 @@ static int doPrintRangeInfo(const CompilerInvocation &InitInvok,
   CompilerInvocation Invocation(InitInvok);
   Invocation.addInputFilename(SourceFileName);
   Invocation.getLangOptions().DisableAvailabilityChecking = false;
+  Invocation.getLangOptions().KeepTokensInSourceFile = true;
 
   CompilerInstance CI;