[SourceKit] Fix SyntaxModel assertion failure due to consuming tokens twice or out of order

lib/IDE/SyntaxModel.cpp

@@ -670,8 +670,13 @@ std::pair<bool, Expr *> ModelASTWalker::walkToExprPre(Expr *E) {
     // from the `VarDecl` and the `PatternBindingDecl` entries.
     // We take over visitation here to avoid walking the `PatternBindingDecl` ones.
     for (auto c : CLE->getCaptureList()) {
-      if (auto *VD = c.Var)
+      if (auto *VD = c.Var) {
+        // We're skipping over the PatternBindingDecl so we need to handle the
+        // the VarDecl's attributes that we'd normally process visiting the PBD.
+        if (!handleAttrs(VD->getAttrs()))
+          return { false, nullptr };
         VD->walk(*this);
+      }
     }
     if (auto *CE = CLE->getClosureBody())
       CE->walk(*this);
@@ -862,8 +867,14 @@ bool ModelASTWalker::walkToDeclPre(Decl *D) {
   if (D->isImplicit())
     return false;
 
-  if (!handleAttrs(D->getAttrs()))
-    return false;
+  // The attributes of EnumElementDecls and VarDecls are handled when visiting
+  // their parent EnumCaseDecl/PatternBindingDecl (which the attributes are
+  // attached to syntactically).
+  if (!isa<EnumElementDecl>(D) &&
+      !(isa<VarDecl>(D) && cast<VarDecl>(D)->getParentPatternBinding())) {
+    if (!handleAttrs(D->getAttrs()))
+      return false;
+  }
 
   if (isa<AccessorDecl>(D)) {
     // Don't push structure nodes for accessors.
@@ -948,6 +959,21 @@ bool ModelASTWalker::walkToDeclPre(Decl *D) {
     SN.TypeRange = charSourceRangeFromSourceRange(SM,
                                       PD->getTypeSourceRangeForDiagnostics());
     pushStructureNode(SN, PD);
+  } else if (auto *PBD = dyn_cast<PatternBindingDecl>(D)) {
+    // Process the attributes of one of the contained VarDecls. Attributes that
+    // are syntactically attached to the PatternBindingDecl end up on the
+    // contained VarDecls.
+    VarDecl *Contained = nullptr;
+    for (auto idx : range(PBD->getNumPatternEntries())) {
+      PBD->getPattern(idx)->forEachVariable([&](VarDecl *VD) -> void {
+        Contained = VD;
+      });
+      if (Contained) {
+        if (!handleAttrs(Contained->getAttrs()))
+          return false;
+        break;
+      }
+    }
   } else if (auto *VD = dyn_cast<VarDecl>(D)) {
     const DeclContext *DC = VD->getDeclContext();
     SyntaxStructureNode SN;
@@ -1012,15 +1038,9 @@ bool ModelASTWalker::walkToDeclPre(Decl *D) {
 
     // We need to handle the special case where attributes semantically
     // attach to enum element decls while syntactically locate before enum case decl.
-    for (auto *EnumElemD : EnumCaseD->getElements()) {
-      for (auto *Att : EnumElemD->getAttrs()) {
-        if (Att->isDeclModifier() &&
-            SM.isBeforeInBuffer(Att->getLocation(), D->getSourceRange().Start)) {
-          passNonTokenNode({SyntaxNodeKind::AttributeBuiltin,
-                            charSourceRangeFromSourceRange(SM,
-                                                           Att->getLocation())});
-        }
-      }
+    if (!EnumCaseD->getElements().empty()) {
+      if (!handleAttrs(EnumCaseD->getElements().front()->getAttrs()))
+        return false;
     }
     if (pushStructureNode(SN, D)) {
       // FIXME: ASTWalker walks enum elements as members of the enum decl, not

test/IDE/coloring.swift

@@ -494,3 +494,31 @@ func typeAttr3(a: @ escaping () -> Int) {}
 
 // CHECK: <kw>func</kw> typeAttr2(a: @ <comment-block>/*this is fine...*/</comment-block> escaping () -> <type>Int</type>, b: <attr-builtin>@ escaping</attr-builtin> () -> <type>Int</type>) {}
 func typeAttr2(a: @ /*this is fine...*/ escaping () -> Int, b: @ escaping () -> Int) {}
+
+// CHECK: <attr-builtin>@available</attr-builtin>(<kw>iOS</kw> <int>99</int>, *)
+// CHECK: <kw>var</kw> iHave = <int>10</int>, multipleVars = <int>20</int>
+@available(iOS 99, *)
+var iHave = 10, multipleVars = 20
+
+enum MultipleCaseElements {
+  // CHECK: <attr-builtin>@available</attr-builtin>(<kw>iOS</kw> <int>99</int>, *)
+  // CHECK: <kw>case</kw> foo, bar
+  @available(iOS 99, *)
+  case foo, bar
+}
+
+protocol P {}
+enum E {
+  // CHECK: <attr-builtin>@available</attr-builtin>(<kw>iOS</kw> <int>99</int>, *)
+  // CHECK: <kw>case</kw> a(<type>P</type>)
+  @available(iOS 99, *)
+  case a(P)
+}
+
+// Ideally this would be attr-builtin, but we don't actually have the attribute
+// in the AST at all.
+//
+// CHECK: <attr-id>@available</attr-id>(<kw>iOS</kw> <int>99</int>, *)
+// CHECK: <kw>var</kw> <kw>_</kw> = <int>10</int>
+@available(iOS 99, *)
+var _ = 10