[Parse] Improve recovery from and diagnostics for invalid operator names

Author: AnthonyLatsis

include/swift/AST/DiagnosticsParse.def

@@ -428,8 +428,14 @@ ERROR(operator_decl_inner_scope,none,
       "'operator' may only be declared at file scope", ())
 ERROR(expected_operator_name_after_operator,PointsToFirstBadToken,
       "expected operator name in operator declaration", ())
-ERROR(identifier_when_expecting_operator,PointsToFirstBadToken,
-      "%0 is considered to be an identifier, not an operator", (Identifier))
+ERROR(identifier_within_operator_name,PointsToFirstBadToken,
+      "'%0' is considered an identifier and must not "
+      "appear within an operator name", (StringRef))
+ERROR(operator_name_invalid_char,PointsToFirstBadToken,
+      "'%0' is not allowed in operator names", (StringRef))
+ERROR(postfix_operator_name_cannot_start_with_unwrap,PointsToFirstBadToken,
+      "postfix operator names starting with '?' or '!' are disallowed "
+      "to avoid collisions with built-in unwrapping operators", ())
 
 ERROR(deprecated_operator_body,PointsToFirstBadToken,
       "operator should no longer be declared with body", ())

lib/Parse/ParseDecl.cpp

@@ -7401,24 +7401,50 @@ Parser::parseDeclOperator(ParseDeclOptions Flags, DeclAttributes &Attributes) {
   SourceLoc OperatorLoc = consumeToken(tok::kw_operator);
   bool AllowTopLevel = Flags.contains(PD_AllowTopLevel);
 
-  if (!Tok.isAnyOperator() && !Tok.is(tok::exclaim_postfix)) {
-    // A common error is to try to define an operator with something in the
-    // unicode plane considered to be an operator, or to try to define an
-    // operator like "not".  Diagnose this specifically.
-    if (Tok.is(tok::identifier))
-      diagnose(Tok, diag::identifier_when_expecting_operator,
-               Context.getIdentifier(Tok.getText()));
-    else
+  const auto maybeDiagnoseInvalidCharInOperatorName = [this](const Token &Tk) {
+    if (Tk.is(tok::identifier) &&
+        DeclAttribute::getAttrKindFromString(Tk.getText()) ==
+          DeclAttrKind::DAK_Count) {
+      diagnose(Tk, diag::identifier_within_operator_name, Tk.getText());
+      return true;
+    } else if (Tk.isNot(tok::colon, tok::l_brace, tok::semi) &&
+               Tk.isPunctuation()) {
+      diagnose(Tk, diag::operator_name_invalid_char,
+               Tk.getText().take_front());
+      return true;
+    }
+    return false;
+  };
+
+  // A common error is to try to define an operator with something in the
+  // unicode plane considered to be an operator, or to try to define an
+  // operator like "not".  Analyze and diagnose this specifically.
+  if (Tok.isAnyOperator()) {
+    if (peekToken().getLoc() == Tok.getRange().getEnd() &&
+      maybeDiagnoseInvalidCharInOperatorName(peekToken())) {
+      consumeToken();
+
+      // If there's a deprecated body, skip it to improve recovery.
+      if (peekToken().is(tok::l_brace)) {
+        consumeToken();
+        skipSingle();
+      }
+      return nullptr;
+    }
+  // We will handle these either below or in Sema.
+  } else if (Tok.isNot(tok::exclaim_postfix, tok::question_infix,
+                       tok::question_postfix)) {
+    if (maybeDiagnoseInvalidCharInOperatorName(Tok)) {
+      // We're done diagnosing.
+    } else {
       diagnose(Tok, diag::expected_operator_name_after_operator);
+    }
 
-    // To improve recovery, check to see if we have a { right after this token.
-    // If so, swallow until the end } to avoid tripping over the body of the
-    // malformed operator decl.
+    // If there's a deprecated body, skip it to improve recovery.
     if (peekToken().is(tok::l_brace)) {
       consumeToken();
       skipSingle();
     }
-
     return nullptr;
   }
 
@@ -7427,9 +7453,11 @@ Parser::parseDeclOperator(ParseDeclOptions Flags, DeclAttributes &Attributes) {
   Identifier Name = Context.getIdentifier(Tok.getText());
   SourceLoc NameLoc = consumeToken();
 
+  // Postfix operators starting with ? or ! conflict with builtin
+  // unwrapping operators.
   if (Attributes.hasAttribute<PostfixAttr>()) {
     if (!Name.empty() && (Name.get()[0] == '?' || Name.get()[0] == '!'))
-      diagnose(NameLoc, diag::expected_operator_name_after_operator);      
+      diagnose(NameLoc, diag::postfix_operator_name_cannot_start_with_unwrap);
   }
 
   auto Result = parseDeclOperatorImpl(OperatorLoc, Name, NameLoc, Attributes);

lib/Sema/TypeCheckAttr.cpp

@@ -1503,6 +1503,9 @@ void AttributeChecker::visitFinalAttr(FinalAttr *attr) {
 static bool isBuiltinOperator(StringRef name, DeclAttribute *attr) {
   return ((isa<PrefixAttr>(attr)  && name == "&") ||   // lvalue to inout
           (isa<PostfixAttr>(attr) && name == "!") ||   // optional unwrapping
+          // FIXME: Not actually a builtin operator, but should probably
+          // be allowed and accounted for in Sema?
+          (isa<PrefixAttr>(attr) && name == "?") ||
           (isa<PostfixAttr>(attr) && name == "?") ||   // optional chaining
           (isa<InfixAttr>(attr) && name == "?") ||     // ternary operator
           (isa<PostfixAttr>(attr) && name == ">") ||   // generic argument list

test/Parse/operator_decl.swift

@@ -40,9 +40,16 @@ prefix operator // expected-error {{expected operator name in operator declarati
 prefix operator %%+
 
 prefix operator ??
-postfix operator ?? // expected-error {{expected operator name in operator declaration}}
+postfix operator ?? // expected-error {{postfix operator names starting with '?' or '!' are disallowed}}
 prefix operator !!
-postfix operator !! // expected-error {{expected operator name in operator declaration}}
+postfix operator !! // expected-error {{postfix operator names starting with '?' or '!' are disallowed}}
+
+infix operator --aa // expected-error {{'aa' is considered an identifier and must not appear within an operator name}}
+infix operator aa--: A // expected-error {{'aa' is considered an identifier and must not appear within an operator name}}
+infix operator <<$$@< // expected-error {{'$$' is considered an identifier and must not appear within an operator name}}
+infix operator !!@aa // expected-error {{'@' is not allowed in operator names}}
+infix operator #++= // expected-error {{'#' is not allowed in operator names}}
+infix operator ++=# // expected-error {{'#' is not allowed in operator names}}
 
 infix operator +++=
 infix operator *** : A

test/Parse/recovery.swift

@@ -818,10 +818,13 @@ func test23719432() {
 }
 
 // <rdar://problem/19911096> QoI: terrible recovery when using '·' for an operator
-infix operator · {  // expected-error {{'·' is considered to be an identifier, not an operator}}
+infix operator · {  // expected-error {{'·' is considered an identifier and must not appear within an operator name}}
   associativity none precedence 150
 }
 
+infix operator -@-class Recover1 {} // expected-error {{'@' is not allowed in operator names}}
+prefix operator -фф--class Recover2 {} // expected-error {{'фф' is considered an identifier and must not appear within an operator name}}
+
 // <rdar://problem/21712891> Swift Compiler bug: String subscripts with range should require closing bracket.
 func r21712891(s : String) -> String {
   let a = s.startIndex..<s.startIndex

test/decl/func/operator.swift

@@ -172,7 +172,7 @@ infix prefix func +-+(x: Double, y: Double) {} // expected-error {{'infix' modif
 
 // Don't allow one to define a postfix '!'; it's built into the
 // language. Also illegal to have any postfix operator starting with '!'.
-postfix operator !  // expected-error {{cannot declare a custom postfix '!' operator}} expected-error {{expected operator name in operator declaration}}
+postfix operator !  // expected-error {{cannot declare a custom postfix '!' operator}} expected-error {{postfix operator names starting with '?' or '!' are disallowed}}
 prefix operator &  // expected-error {{cannot declare a custom prefix '&' operator}}
 
 // <rdar://problem/14607026> Restrict use of '<' and '>' as prefix/postfix operator names
@@ -188,7 +188,9 @@ func operator_in_func_bad () {
     prefix func + (input: String) -> String { return "+" + input } // expected-error {{operator functions can only be declared at global or in type scope}}
 }
 
-infix operator ?  // expected-error {{expected operator name in operator declaration}} 
+infix operator ?  // expected-error {{cannot declare a custom infix '?' operator}}
+
+prefix operator ?  // expected-error {{cannot declare a custom prefix '?' operator}}
 
 infix operator ??=