Fix parsing of @_unavailableFromAsync

There’s a very easy to reach `llvm_unreachable()` in this code which ought to be a diagnostic, as well as a couple of other issues. Rework it into something that’s a bit better at handling the edge cases.

Author: beccadax

include/swift/AST/KnownIdentifiers.def

@@ -124,6 +124,7 @@ IDENTIFIER(AsyncIterator)
 IDENTIFIER(load)
 IDENTIFIER(main)
 IDENTIFIER_WITH_NAME(MainEntryPoint, "$main")
+IDENTIFIER(message)
 IDENTIFIER(next)
 IDENTIFIER_(nsErrorDomain)
 IDENTIFIER(objectAtIndexedSubscript)

include/swift/Parse/Parser.h

@@ -1782,7 +1782,10 @@ class Parser {
   /// \param name The parsed name of the label (empty if it doesn't exist, or is
   /// _)
   /// \param loc The location of the label (empty if it doesn't exist)
-  void parseOptionalArgumentLabel(Identifier &name, SourceLoc &loc);
+  /// \param isAttr True if this is an argument label for an attribute (allows, but deprecates, use of
+  ///               \c '=' instead of \c ':').
+  void parseOptionalArgumentLabel(Identifier &name, SourceLoc &loc,
+                                  bool isAttr = false);
 
   enum class DeclNameFlag : uint8_t {
     /// If passed, operator basenames are allowed.

lib/Parse/ParseDecl.cpp

@@ -3780,35 +3780,60 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
   case DeclAttrKind::UnavailableFromAsync: {
     StringRef message;
     if (consumeIfAttributeLParen()) {
-      if (!Tok.is(tok::identifier)) {
-        llvm_unreachable("Flag must start with an identifier");
-      }
-
-      StringRef flag = Tok.getText();
+      auto tokMayBeArgument = [&]() -> bool {
+        return Tok.isNot(tok::r_paren, tok::comma) &&
+                 !isKeywordPossibleDeclStart(Context.LangOpts, Tok);
+      };
 
-      if (flag != "message") {
-        diagnose(Tok.getLoc(), diag::attr_unknown_option, flag, AttrName);
-        return makeParserError();
-      }
-      consumeToken();
-      if (!consumeIf(tok::colon)) {
-        if (!Tok.is(tok::equal)) {
-          diagnose(Tok.getLoc(), diag::attr_expected_colon_after_label, flag);
-          return makeParserSuccess();
+      Identifier label;
+      SourceLoc labelLoc;
+      parseOptionalArgumentLabel(label, labelLoc, /*isAttr=*/true);
+
+      if (label.empty()) {
+        // If we have the identifier 'message', assume the user forgot the
+        // colon.
+        if (Tok.isContextualKeyword("message")) {
+          labelLoc = consumeToken();
+          auto diag = diagnose(Tok, diag::attr_expected_colon_after_label,
+                               "message");
+          if (Tok.is(tok::string_literal))
+            diag.fixItInsertAfter(labelLoc, ":");
+          else
+            return makeParserError();
         }
-        diagnose(Tok.getLoc(), diag::replace_equal_with_colon_for_value)
-          .fixItReplace(Tok.getLoc(), ": ");
-        consumeToken();
+        // If the argument list just abruptly cuts off, handle that as a
+        // missing argument (below). Otherwise, diagnose the missing label.
+        else if (tokMayBeArgument()) {
+          if (labelLoc.isValid())
+            // The user wrote an explicitly omitted label (`_:`).
+            diagnose(labelLoc, diag::attr_expected_label, "message", AttrName)
+              .fixItReplace(labelLoc, "message");
+          else
+            diagnose(Tok, diag::attr_expected_label, "message", AttrName)
+              .fixItInsert(Tok.getLoc(), "message: ");
+        }
+        // Fall through to parse the argument.
+      } else if (label != Context.Id_message) {
+        diagnose(labelLoc, diag::attr_unknown_option, label.str(), AttrName)
+          .fixItReplace(labelLoc, "message");
+        return makeParserError();
       }
+
       if (!Tok.is(tok::string_literal)) {
-        diagnose(Tok.getLoc(), diag::attr_expected_string_literal, AttrName);
-        return makeParserSuccess();
+        // If this token looks like an argument, replace it; otherwise insert
+        // before it.
+        auto endLoc = tokMayBeArgument() ? peekToken().getLoc() : Tok.getLoc();
+
+        diagnose(Tok, diag::attr_expected_string_literal, AttrName)
+          .fixItReplaceChars(Tok.getLoc(), endLoc, "\"<#error message#>\"");
+
+        return makeParserError();
       }
 
       std::optional<StringRef> value =
-          getStringLiteralIfNotInterpolated(Tok.getLoc(), flag);
+          getStringLiteralIfNotInterpolated(Tok.getLoc(), "message");
       if (!value)
-        return makeParserSuccess();
+        return makeParserError();
       Token stringTok = Tok;
       consumeToken();
       message = *value;

lib/Parse/ParseExpr.cpp

@@ -2129,9 +2129,14 @@ ParserResult<Expr> Parser::parseExprStringLiteral() {
                                       AppendingExpr));
 }
 
-void Parser::parseOptionalArgumentLabel(Identifier &name, SourceLoc &loc) {
+void Parser::parseOptionalArgumentLabel(Identifier &name, SourceLoc &loc,
+                                        bool isAttr) {
+  /// A token that has the same meaning as colon, but is deprecated, if one exists for this call.
+  auto altColon = isAttr ? tok::equal : tok::NUM_TOKENS;
+
   // Check to see if there is an argument label.
-  if (Tok.canBeArgumentLabel() && peekToken().is(tok::colon)) {
+  if (Tok.canBeArgumentLabel() && peekToken().isAny(tok::colon, altColon)) {
+    // Label found, including colon.
     auto text = Tok.getText();
 
     // If this was an escaped identifier that need not have been escaped, say
@@ -2149,11 +2154,23 @@ void Parser::parseOptionalArgumentLabel(Identifier &name, SourceLoc &loc) {
     }
 
     loc = consumeArgumentLabel(name, /*diagnoseDollarPrefix=*/false);
-    consumeToken(tok::colon);
-  } else if (Tok.is(tok::colon)) {
-    diagnose(Tok, diag::expected_label_before_colon);
-    consumeToken(tok::colon);
+  } else if (Tok.isAny(tok::colon, altColon)) {
+    // Found only the colon.
+    diagnose(Tok, diag::expected_label_before_colon)
+      .fixItInsert(Tok.getLoc(), "<#label#>");
+  } else {
+    // No label here.
+    return;
   }
+
+  // If we get here, we ought to be on the colon.
+  assert(Tok.isAny(tok::colon, altColon));
+
+  if (Tok.is(altColon))
+    diagnose(Tok, diag::replace_equal_with_colon_for_value)
+      .fixItReplace(Tok.getLoc(), ": ");
+
+  consumeToken();
 }
 
 static bool tryParseArgLabelList(Parser &P, Parser::DeclNameOptions flags,

test/Concurrency/unavailable_from_async.swift

@@ -140,25 +140,54 @@ func asyncFunc() async { // expected-error{{asynchronous global function 'asyncF
 
 // Parsing tests
 
+// expected-error@+1 {{expected label 'message:' in '@_unavailableFromAsync' attribute}} {{24-24=message: }}
+@_unavailableFromAsync("almost right, but not quite")
+func blarp0() {}
+
+// expected-error@+1 {{expected label 'message:' in '@_unavailableFromAsync' attribute}} {{24-25=message}}
+@_unavailableFromAsync(_: "almost right, but not quite")
+func blarp0a() {}
+
 // expected-error@+2 {{expected declaration}}
-// expected-error@+1:24{{unknown option 'nope' for attribute '_unavailableFromAsync'}}
+// expected-error@+1:24{{unknown option 'nope' for attribute '_unavailableFromAsync'}} {{24-28=message}}
 @_unavailableFromAsync(nope: "almost right, but not quite")
 func blarp1() {}
 
 // expected-error@+2 {{expected declaration}}
-// expected-error@+1 {{expected ':' after label 'message'}}
+// expected-error@+1 {{expected ':' after label 'message'}} {{none}}
 @_unavailableFromAsync(message; "almost right, but not quite")
 func blarp2() {}
 
+// expected-error@+1 {{expected ':' after label 'message'}} {{31-31=:}}
+@_unavailableFromAsync(message "almost right, but not quite")
+func blarp2a() {}
+
 // expected-error@+1:31 {{'=' has been replaced with ':' in attribute arguments}}{{31-32=: }}
 @_unavailableFromAsync(message="almost right, but not quite")
 func blarp3() {}
 
 // expected-error@+2 {{expected declaration}}
-// expected-error@+1 {{expected string literal in '_unavailableFromAsync' attribute}}
+// expected-error@+1 {{expected string literal in '_unavailableFromAsync' attribute}} {{33-35="<#error message#>"}}
 @_unavailableFromAsync(message: 32)
 func blarp4() {}
 
+// expected-error@+2 {{expected declaration}}
+// expected-error@+1:24{{unknown option 'fnord' for attribute '_unavailableFromAsync'}} {{24-29=message}}
+@_unavailableFromAsync(fnord: 32)
+func blarp4a() {}
+
+// expected-error@+3 {{expected declaration}}
+// expected-error@+2 {{expected label 'message:' in '@_unavailableFromAsync' attribute}} {{24-25=message}}
+// expected-error@+1 {{expected string literal in '_unavailableFromAsync' attribute}} {{27-29="<#error message#>"}}
+@_unavailableFromAsync(_: 32)
+func blarp4b() {}
+
+// expected-error@+3 {{expected declaration}}
+// expected-error@+2 {{expected string literal in '_unavailableFromAsync' attribute}} {{24-26="<#error message#>"}}
+// expected-error@+1 {{expected label 'message:' in '@_unavailableFromAsync' attribute}} {{24-24=message: }}
+@_unavailableFromAsync(32)
+func blarp4c() {}
+
 // expected-error@+2 {{expected declaration}}
 // expected-error@+1 {{message cannot be an interpolated string}}
 @_unavailableFromAsync(message: "blarppy blarp \(31 + 10)")