SourceKit/SyntaxMap: properly handle object literal syntax for highlighting

We should use parser to figure out the end position of object literal expression instead
of scanning through token stream, which crashes sourcekitd when the syntax is invalid.

Fixing: rdar://48390913

Author: nkcsgexi

lib/IDE/SyntaxModel.cpp

@@ -56,7 +56,6 @@ SyntaxModelContext::SyntaxModelContext(SourceFile &SrcFile)
   std::vector<SyntaxNode> Nodes;
   SourceLoc AttrLoc;
   SourceLoc UnaryMinusLoc;
-  auto LiteralStartLoc = Optional<SourceLoc>();
   for (unsigned I = 0, E = Tokens.size(); I != E; ++I) {
     auto &Tok = Tokens[I];
     // Ignore empty string literals between interpolations, e.g. "\(1)\(2)"
@@ -85,16 +84,6 @@ SyntaxModelContext::SyntaxModelContext(SourceFile &SrcFile)
       Loc = Tok.getLoc();
       Length = Tok.getLength();
 
-      if (LiteralStartLoc.hasValue() && Length.hasValue()) {
-        if (Tok.getKind() != tok::r_paren)
-          continue;
-        Kind = SyntaxNodeKind::ObjectLiteral;
-        Nodes.emplace_back(Kind, CharSourceRange(SM, LiteralStartLoc.getValue(),
-                                                 Tok.getRange().getEnd()));
-        LiteralStartLoc = Optional<SourceLoc>();
-        continue;
-      }
-
       switch(Tok.getKind()) {
 #define KEYWORD(X) case tok::kw_##X:
 #include "swift/Syntax/TokenKinds.def"
@@ -105,8 +94,8 @@ SyntaxModelContext::SyntaxModelContext(SourceFile &SrcFile)
 
 #define POUND_OBJECT_LITERAL(Name, Desc, Proto) case tok::pound_##Name:
 #include "swift/Syntax/TokenKinds.def"
-        LiteralStartLoc = Loc;
-        continue;
+        Kind = SyntaxNodeKind::Keyword;
+        break;
 
 #define POUND_COND_DIRECTIVE_KEYWORD(Name) case tok::pound_##Name:
 #include "swift/Syntax/TokenKinds.def"
@@ -260,6 +249,9 @@ class ModelASTWalker : public ASTWalker {
   Optional<SyntaxNode> parseFieldNode(StringRef Text, StringRef OrigText,
                                       SourceLoc OrigLoc);
   llvm::DenseSet<ASTNode> VisitedNodesInsideIfConfig;
+  /// When non-zero, we should avoid passing tokens as syntax nodes since a parent of several tokens
+  /// is considered as one, e.g. object literal expression.
+  uint8_t AvoidPassingSyntaxToken = 0;
 
 public:
   SyntaxModelWalker &Walker;
@@ -494,8 +486,9 @@ std::pair<bool, Expr *> ModelASTWalker::walkToExprPre(Expr *E) {
     SN.NameRange = CharSourceRange(SM, NRStart, NREnd);
     SN.BodyRange =
       innerCharSourceRangeFromSourceRange(SM, ObjectE->getSourceRange());
+    // Consider the object literal as a single syntax token for highlighting.
+    passNonTokenNode({SyntaxNodeKind::ObjectLiteral, SN.Range});
     pushStructureNode(SN, E);
-
   } else if (auto *ArrayE = dyn_cast<ArrayExpr>(E)) {
     SyntaxStructureNode SN;
     SN.Kind = SyntaxStructureKind::ArrayExpression;
@@ -1152,8 +1145,10 @@ bool ModelASTWalker::passTokenNodesUntil(SourceLoc Loc,
       }
       break;
     }
-    if (!passNode(TokenNodes[I]))
-      return false;
+    if (!AvoidPassingSyntaxToken) {
+      if (!passNode(TokenNodes[I]))
+        return false;
+    }
   }
 
   TokenNodes = TokenNodes.slice(I);
@@ -1200,9 +1195,15 @@ bool ModelASTWalker::passNode(const SyntaxNode &Node) {
   return Walker.walkToNodePost(Node);
 }
 
+static bool shouldAvoidPssingSyntaxToken(const SyntaxStructureNode &Node) {
+  return Node.Kind == SyntaxStructureKind::ObjectLiteralExpression;
+}
+
 bool ModelASTWalker::pushStructureNode(const SyntaxStructureNode &Node,
                                        const ASTNodeType& ASTNode) {
   SubStructureStack.emplace_back(Node, ASTNode);
+  if (shouldAvoidPssingSyntaxToken(Node))
+    AvoidPassingSyntaxToken ++;
 
   if (!passTokenNodesUntil(Node.Range.getStart(), ExcludeNodeAtLocation))
     return false;
@@ -1215,6 +1216,12 @@ bool ModelASTWalker::pushStructureNode(const SyntaxStructureNode &Node,
 bool ModelASTWalker::popStructureNode() {
   assert(!SubStructureStack.empty());
   SyntaxStructureNode Node = SubStructureStack.back().StructureNode;
+  SWIFT_DEFER {
+    if (shouldAvoidPssingSyntaxToken(Node)) {
+      assert(AvoidPassingSyntaxToken);
+      AvoidPassingSyntaxToken --;
+    }
+  };
   SubStructureStack.pop_back();
 
   // VarDecls are popped before we see their TypeRepr, so if we pass the token

test/IDE/coloring.swift

@@ -262,6 +262,9 @@ func bar(x: Int) -> (Int, Float) {
   foo(Float())
 }
 
+// CHECK: <object-literal>#colorLiteral(red: 1.0, green: 1.0, blue: 1.0, alpha: 1.0)</object-literal>
+#colorLiteral(red: 1.0, green: 1.0, blue: 1.0, alpha: 1.0)
+
 class GenC<T1,T2> {}
 
 func test() {
@@ -277,6 +280,9 @@ func test2(x: Int) {
   "\(x)"
 }
 
+// CHECK: <kw>#colorLiteral</kw>
+#colorLiteral
+
 // CHECK: <kw>class</kw> Observers {
 class Observers {
   // CHECK: <kw>var</kw> p1 : <type>Int</type> {