[syntax-map] Fix array-of-object-literal syntax map, argument label keywords

Argument labels are allowed to use keywords, in which case we want to
treat them as identifiers in the syntax map (except for '_'). This
commit moves calculation of that into the original lexing instead of
in the model walker, which makes it much more robust, since the model
walker was only guessing about what was next on the the TokenNodes list.

This fixes a bug where arrays of object literals would only have the
first object correct (the following ones were identifiers), as well as
some incorrect cases where we treated keywords as identifiers.

rdar://problem/27726422

Author: benlangmuir

lib/IDE/SyntaxModel.cpp

@@ -94,9 +94,23 @@ SyntaxModelContext::SyntaxModelContext(SourceFile &SrcFile)
       }
 
       switch(Tok.getKind()) {
-#define KEYWORD(X) case tok::kw_##X: Kind = SyntaxNodeKind::Keyword; break;
+#define KEYWORD(X) case tok::kw_##X:
 #include "swift/Parse/Tokens.def"
 #undef KEYWORD
+        if (Tok.getKind() != tok::kw__ &&
+            0 < I && I < Tokens.size() - 1 &&
+            (Tokens[I-1].getKind() == tok::l_paren ||
+             Tokens[I-1].getKind() == tok::comma) &&
+            (Tokens[I+1].getKind() == tok::colon ||
+             Tokens[I+1].getKind() == tok::identifier ||
+             Tokens[I+1].isKeyword())) {
+          // Keywords are allowed as argument labels and should be treated as
+          // identifiers.  The exception is '_' which is not a name.
+          Kind = SyntaxNodeKind::Identifier;
+        } else {
+          Kind = SyntaxNodeKind::Keyword;
+        }
+        break;
 
 #define POUND_OLD_OBJECT_LITERAL(Name, NewName, OldArg, NewArg) \
       case tok::pound_##Name:
@@ -453,9 +467,6 @@ std::pair<bool, Expr *> ModelASTWalker::walkToExprPre(Expr *E) {
                                                                   E->getEndLoc()));
         passTokenNodesUntil(NR.getStart(),
                             PassNodesBehavior::ExcludeNodeAtLocation);
-        if (!TokenNodes.empty())
-          const_cast<SyntaxNode&>(TokenNodes.front()).Kind = SyntaxNodeKind::
-            Identifier;
       }
       else
         SN.Range = SN.BodyRange;
@@ -524,13 +535,8 @@ std::pair<bool, Expr *> ModelASTWalker::walkToExprPre(Expr *E) {
   } else if (auto *Tup = dyn_cast<TupleExpr>(E)) {
     for (unsigned I = 0; I < Tup->getNumElements(); ++ I) {
       SourceLoc NameLoc = Tup->getElementNameLoc(I);
-      if (NameLoc.isValid()) {
+      if (NameLoc.isValid())
         passTokenNodesUntil(NameLoc, PassNodesBehavior::ExcludeNodeAtLocation);
-        if (!TokenNodes.empty()) {
-          const_cast<SyntaxNode&>(TokenNodes.front()).Kind = SyntaxNodeKind::
-            Identifier;
-        }
-      }
     }
   }
 
@@ -860,8 +866,6 @@ bool ModelASTWalker::walkToDeclPre(Decl *D) {
       SourceLoc ArgStart = PD->getSourceRange().Start;
       SN.NameRange = CharSourceRange(ArgStart, PD->getArgumentName().getLength());
       passTokenNodesUntil(ArgStart, PassNodesBehavior::ExcludeNodeAtLocation);
-      const_cast<SyntaxNode&>(TokenNodes.front()).Kind = SyntaxNodeKind::
-        Identifier;
     }
     SN.Range = charSourceRangeFromSourceRange(SM, PD->getSourceRange());
     SN.Attrs = PD->getAttrs();

test/IDE/coloring.swift

@@ -273,7 +273,7 @@ func test3(o: AnyObject) {
   let x = o as! MyCls
 }
 
-// CHECK: <kw>func</kw> test4(<kw>inout</kw> a: <type>Int</type>) {{{$}}
+// CHECK: <kw>func</kw> test4(inout a: <type>Int</type>) {{{$}}
 func test4(inout a: Int) {
   // CHECK: <kw>if</kw> <kw>#available</kw> (<kw>OSX</kw> >= <float>10.10</float>, <kw>iOS</kw> >= <float>8.01</float>) {<kw>let</kw> OSX = <str>"iOS"</str>}}{{$}}
   if #available (OSX >= 10.10, iOS >= 8.01) {let OSX = "iOS"}}
@@ -501,20 +501,44 @@ let file = #fileLiteral(resourceName: "cloud.png")
 let black = #colorLiteral(red: 0, green: 0, blue: 0, alpha: 1)
 // CHECK: <kw>let</kw> black = <object-literal>#colorLiteral(red: 0, green: 0, blue: 0, alpha: 1)</object-literal>
 
+let rgb = [#colorLiteral(red: 1, green: 0, blue: 0, alpha: 1),
+           #colorLiteral(red: 0, green: 1, blue: 0, alpha: 1),
+           #colorLiteral(red: 0, green: 0, blue: 1, alpha: 1)]
+// CHECK: <kw>let</kw> rgb = [<object-literal>#colorLiteral(red: 1, green: 0, blue: 0, alpha: 1)</object-literal>,
+// CHECK:                     <object-literal>#colorLiteral(red: 0, green: 1, blue: 0, alpha: 1)</object-literal>,
+// CHECK:                     <object-literal>#colorLiteral(red: 0, green: 0, blue: 1, alpha: 1)</object-literal>]
+
 "--\"\(x) --"
 // CHECK: <str>"--\"</str>\<anchor>(</anchor>x<anchor>)</anchor><str> --"</str>
 
 func keywordAsLabel1(in: Int) {}
 // CHECK: <kw>func</kw> keywordAsLabel1(in: <type>Int</type>) {}
 func keywordAsLabel2(for: Int) {}
 // CHECK: <kw>func</kw> keywordAsLabel2(for: <type>Int</type>) {}
+func keywordAsLabel3(if: Int, for: Int) {}
+// CHECK: <kw>func</kw> keywordAsLabel3(if: <type>Int</type>, for: <type>Int</type>) {}
+func keywordAsLabel4(_: Int) {}
+// CHECK: <kw>func</kw> keywordAsLabel4(<kw>_</kw>: <type>Int</type>) {}
+func keywordAsLabel5(_: Int, for: Int) {}
+// CHECK: <kw>func</kw> keywordAsLabel5(<kw>_</kw>: <type>Int</type>, for: <type>Int</type>) {}
+func keywordAsLabel6(if let: Int) {}
+// CHECK: <kw>func</kw> keywordAsLabel6(if <kw>let</kw>: <type>Int</type>) {}
 
 func foo1() {
 // CHECK: <kw>func</kw> foo1() {
   keywordAsLabel1(in: 1)
 // CHECK: keywordAsLabel1(in: <int>1</int>)
   keywordAsLabel2(for: 1)
 // CHECK: keywordAsLabel2(for: <int>1</int>)
+  keywordAsLabel3(if: 1, for: 2)
+// CHECK: keywordAsLabel3(if: <int>1</int>, for: <int>2</int>)
+  keywordAsLabel5(1, for: 2)
+// CHECK: keywordAsLabel5(<int>1</int>, for: <int>2</int>)
+
+  _ = (if: 0, for: 2)
+// CHECK: <kw>_</kw> = (if: <int>0</int>, for: <int>2</int>)
+  _ = (_: 0, _: 2)
+// CHECK: <kw>_</kw> = (<kw>_</kw>: <int>0</int>, <kw>_</kw>: <int>2</int>)
 }
 
 // Keep this as the last test

test/SourceKit/SyntaxMapData/syntaxmap-object-literals.swift.response

@@ -29,7 +29,7 @@
       key.length: 40
     },
     {
-      key.kind: source.lang.swift.syntaxtype.identifier,
+      key.kind: source.lang.swift.syntaxtype.keyword,
       key.offset: 165,
       key.length: 3
     },
@@ -44,7 +44,7 @@
       key.length: 50
     },
     {
-      key.kind: source.lang.swift.syntaxtype.identifier,
+      key.kind: source.lang.swift.syntaxtype.keyword,
       key.offset: 228,
       key.length: 3
     },