Parser: use parser to generate a refined token stream to help syntax coloring. (#11809)

This patch allows Parser to generate a refined token stream to satisfy tooling's need. For syntax coloring, token stream from lexer is insufficient because (1) we have contextual keywords like get and set; (2) we may allow keywords to be used as argument labels and names; and (3) we need to split tokens like "==<". In this patch, these refinements are directly fulfilled through parsing without additional heuristics. The refined token vector is optionally saved in SourceFile instance.

Author: nkcsgexi

include/swift/AST/Module.h

@@ -27,6 +27,7 @@
 #include "swift/Basic/OptionSet.h"
 #include "swift/Basic/SourceLoc.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/Parse/Token.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
@@ -886,7 +887,7 @@ class SourceFile final : public FileUnit {
   ASTStage_t ASTStage = Parsing;
 
   SourceFile(ModuleDecl &M, SourceFileKind K, Optional<unsigned> bufferID,
-             ImplicitModuleImportKind ModImpKind);
+             ImplicitModuleImportKind ModImpKind, bool KeepTokens);
 
   void
   addImports(ArrayRef<std::pair<ModuleDecl::ImportedModule, ImportOptions>> IM);
@@ -1070,6 +1071,25 @@ class SourceFile final : public FileUnit {
     getInterfaceHash(str);
     out << str << '\n';
   }
+
+  std::vector<Token> &getTokenVector() {
+    assert(EnabledAndAllCorrectedTokens.first && "Disabled");
+    return EnabledAndAllCorrectedTokens.second;
+  }
+
+  ArrayRef<Token> getAllTokens() const {
+    assert(EnabledAndAllCorrectedTokens.first && "Disabled");
+    return EnabledAndAllCorrectedTokens.second;
+  }
+
+  bool shouldKeepTokens() const {
+    return EnabledAndAllCorrectedTokens.first;
+  }
+
+private:
+  /// Whether the SourceFile instance opts to collect underlying tokens and
+  /// the vector containing these tokens if so.
+  std::pair<bool, std::vector<Token>> EnabledAndAllCorrectedTokens;
 };
 
 

include/swift/Basic/LangOptions.h

@@ -245,6 +245,11 @@ namespace swift {
     /// This is used to guard preemptive testing for the fix-it.
     bool FixStringToSubstringConversions = false;
 
+    /// Whether to keep track of a refined token stream in SourceFile while
+    /// parsing. This is set true usually for tooling purposes like semantic
+    /// coloring.
+    bool KeepTokensInSourceFile = false;
+
     /// Sets the target we are building for and updates platform conditions
     /// to match.
     ///

include/swift/Parse/Lexer.h

@@ -411,6 +411,10 @@ class Lexer {
       Result.IndentToStrip = 0;
       return Result;
     }
+
+    SourceLoc getEndLoc() {
+      return Loc.getAdvancedLoc(Length);
+    }
   };
 
   /// \brief Compute the bytes that the actual string literal should codegen to.

include/swift/Parse/Parser.h

@@ -71,6 +71,22 @@ namespace swift {
     ActiveConditionalBlock,
   };
 
+/// The receiver will be fed with consumed tokens while parsing. The main purpose
+/// is to generate a corrected token stream for tooling support like syntax
+/// coloring.
+class ConsumeTokenReceiver {
+public:
+  /// This is called when a token is consumed.
+  virtual void receive(Token Tok) {}
+
+  /// This is called to update the kind of a token whose start location is Loc.
+  virtual void registerTokenKindChange(SourceLoc Loc, tok NewKind) {};
+
+  /// This is called when a source file is fully parsed.
+  virtual void finalize() {};
+  virtual ~ConsumeTokenReceiver() = default;
+};
+
 /// The main class used for parsing a source file (.swift or .sil).
 ///
 /// Rather than instantiating a Parser yourself, use one of the parsing APIs
@@ -174,6 +190,9 @@ class Parser {
   /// \brief This is the current token being considered by the parser.
   Token Tok;
 
+  /// \brief The receiver to collect all consumed tokens.
+  ConsumeTokenReceiver *TokReceiver;
+
   /// \brief The location of the previous token.
   SourceLoc PreviousLoc;
 
@@ -356,7 +375,7 @@ class Parser {
     // We might be at tok::eof now, so ensure that consumeToken() does not
     // assert about lexing past eof.
     Tok.setKind(tok::unknown);
-    consumeToken();
+    consumeTokenWithoutFeedingReceiver();
 
     PreviousLoc = PP.PreviousLoc;
   }
@@ -369,7 +388,7 @@ class Parser {
     // We might be at tok::eof now, so ensure that consumeToken() does not
     // assert about lexing past eof.
     Tok.setKind(tok::unknown);
-    consumeToken();
+    consumeTokenWithoutFeedingReceiver();
 
     PreviousLoc = PP.PreviousLoc;
   }
@@ -421,6 +440,10 @@ class Parser {
   /// \brief Return the next token that will be installed by \c consumeToken.
   const Token &peekToken();
 
+  /// Consume a token that we created on the fly to correct the original token
+  /// stream from lexer.
+  void consumeExtraToken(Token K);
+  SourceLoc consumeTokenWithoutFeedingReceiver();
   SourceLoc consumeToken();
   SourceLoc consumeToken(tok K) {
     assert(Tok.is(K) && "Consuming wrong token kind");

include/swift/Parse/Token.h

@@ -61,13 +61,17 @@ class Token {
 public:
   Token() : Kind(tok::NUM_TOKENS), AtStartOfLine(false), CommentLength(0),
             EscapedIdentifier(false) {}
-  Token(tok Kind, StringRef Text)
-    : Kind(Kind), AtStartOfLine(false), CommentLength(0),
-      EscapedIdentifier(false), MultilineString(false),
-      Text(Text) {}
-  
+
+  Token(tok Kind, StringRef Text, unsigned CommentLength)
+          : Kind(Kind), AtStartOfLine(false), CommentLength(CommentLength),
+            EscapedIdentifier(false), MultilineString(false),
+            Text(Text) {}
+
+  Token(tok Kind, StringRef Text): Token(Kind, Text, 0) {};
+
   tok getKind() const { return Kind; }
   void setKind(tok K) { Kind = K; }
+  void clearCommentLength() { CommentLength = 0; }
 
   /// is/isNot - Predicates to check if this token is a specific kind, as in
   /// "if (Tok.is(tok::l_brace)) {...}".

include/swift/Syntax/TokenKinds.def

@@ -287,6 +287,9 @@ MISC(dollarident)
 MISC(sil_local_name)       // %42 in SIL mode.
 MISC(comment)
 
+MISC(string_interpolation_anchor)
+MISC(contextual_keyword)
+
 #undef KEYWORD
 #undef SWIFT_KEYWORD
 #undef DECL_KEYWORD

lib/AST/Module.cpp

@@ -1321,9 +1321,12 @@ static void performAutoImport(SourceFile &SF,
 
 SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
                        Optional<unsigned> bufferID,
-                       ImplicitModuleImportKind ModImpKind)
+                       ImplicitModuleImportKind ModImpKind,
+                       bool KeepTokens)
   : FileUnit(FileUnitKind::Source, M),
-    BufferID(bufferID ? *bufferID : -1), Kind(K) {
+    BufferID(bufferID ? *bufferID : -1),
+    Kind(K),
+    EnabledAndAllCorrectedTokens(KeepTokens, std::vector<Token>()) {
   M.getASTContext().addDestructorCleanup(*this);
   performAutoImport(*this, ModImpKind);
 

lib/Frontend/Frontend.cpp

@@ -264,6 +264,7 @@ ModuleDecl *CompilerInstance::getMainModule() {
 void CompilerInstance::performSema() {
   const FrontendOptions &options = Invocation.getFrontendOptions();
   const InputFileKind Kind = Invocation.getInputKind();
+  bool KeepTokens = Invocation.getLangOptions().KeepTokensInSourceFile;
   ModuleDecl *MainModule = getMainModule();
   Context->LoadedModules[MainModule->getName()] = MainModule;
 
@@ -392,7 +393,7 @@ void CompilerInstance::performSema() {
   if (Kind == InputFileKind::IFK_Swift_REPL) {
     auto *SingleInputFile =
       new (*Context) SourceFile(*MainModule, Invocation.getSourceFileKind(),
-                                None, modImpKind);
+                                None, modImpKind, KeepTokens);
     MainModule->addFile(*SingleInputFile);
     addAdditionalInitialImports(SingleInputFile);
     return;
@@ -420,7 +421,8 @@ void CompilerInstance::performSema() {
 
     auto *MainFile = new (*Context) SourceFile(*MainModule,
                                                Invocation.getSourceFileKind(),
-                                               MainBufferID, modImpKind);
+                                               MainBufferID, modImpKind,
+                                               KeepTokens);
     MainModule->addFile(*MainFile);
     addAdditionalInitialImports(MainFile);
 
@@ -446,7 +448,8 @@ void CompilerInstance::performSema() {
     auto *NextInput = new (*Context) SourceFile(*MainModule,
                                                 SourceFileKind::Library,
                                                 BufferID,
-                                                modImpKind);
+                                                modImpKind,
+                                                KeepTokens);
     MainModule->addFile(*NextInput);
     addAdditionalInitialImports(NextInput);
 
@@ -585,6 +588,7 @@ void CompilerInstance::performParseOnly(bool EvaluateConditionals) {
   const InputFileKind Kind = Invocation.getInputKind();
   ModuleDecl *MainModule = getMainModule();
   Context->LoadedModules[MainModule->getName()] = MainModule;
+  bool KeepTokens = Invocation.getLangOptions().KeepTokensInSourceFile;
 
   assert((Kind == InputFileKind::IFK_Swift ||
           Kind == InputFileKind::IFK_Swift_Library) &&
@@ -600,7 +604,8 @@ void CompilerInstance::performParseOnly(bool EvaluateConditionals) {
     SourceMgr.setHashbangBufferID(MainBufferID);
 
     auto *MainFile = new (*Context) SourceFile(
-        *MainModule, Invocation.getSourceFileKind(), MainBufferID, modImpKind);
+        *MainModule, Invocation.getSourceFileKind(), MainBufferID, modImpKind,
+                                               KeepTokens);
     MainModule->addFile(*MainFile);
 
     if (MainBufferID == PrimaryBufferID)
@@ -615,7 +620,9 @@ void CompilerInstance::performParseOnly(bool EvaluateConditionals) {
       continue;
 
     auto *NextInput = new (*Context)
-        SourceFile(*MainModule, SourceFileKind::Library, BufferID, modImpKind);
+        SourceFile(*MainModule, SourceFileKind::Library, BufferID, modImpKind,
+                   KeepTokens);
+
     MainModule->addFile(*NextInput);
     if (BufferID == PrimaryBufferID)
       setPrimarySourceFile(NextInput);

lib/IDE/Formatting.cpp

@@ -485,11 +485,11 @@ class FormatContext {
 };
 
 class FormatWalker : public SourceEntityWalker {
-  typedef std::vector<Token>::iterator TokenIt;
+  typedef ArrayRef<Token>::iterator TokenIt;
   class SiblingCollector {
     SourceLoc FoundSibling;
     SourceManager &SM;
-    std::vector<Token> &Tokens;
+    ArrayRef<Token> Tokens;
     SourceLoc &TargetLoc;
     TokenIt TI;
     bool NeedExtraIndentation;
@@ -552,7 +552,7 @@ class FormatWalker : public SourceEntityWalker {
     }
 
   public:
-    SiblingCollector(SourceManager &SM, std::vector<Token> &Tokens,
+    SiblingCollector(SourceManager &SM, ArrayRef<Token> Tokens,
                      SourceLoc &TargetLoc) : SM(SM), Tokens(Tokens),
     TargetLoc(TargetLoc), TI(Tokens.begin()),
     NeedExtraIndentation(false) {}
@@ -660,7 +660,7 @@ class FormatWalker : public SourceEntityWalker {
   bool InDocCommentBlock = false;
   bool InCommentLine = false;
   bool InStringLiteral = false;
-  std::vector<Token> Tokens;
+  ArrayRef<Token> Tokens;
   LangOptions Options;
   TokenIt CurrentTokIt;
   unsigned TargetLine;
@@ -739,7 +739,7 @@ class FormatWalker : public SourceEntityWalker {
 public:
   explicit FormatWalker(SourceFile &SF, SourceManager &SM)
   :SF(SF), SM(SM),
-  Tokens(tokenize(Options, SM, SF.getBufferID().getValue())),
+  Tokens(SF.getAllTokens()),
   CurrentTokIt(Tokens.begin()),
   SCollector(SM, Tokens, TargetLocation) {}