fix @differentiable(linear) type attr parsing (#27669)

Author: marcrasi

lib/Parse/ASTGen.cpp

@@ -451,6 +451,15 @@ TypeAttributes ASTGen::generateTypeAttributes(const AttributeListSyntax &syntax,
       if (indexTok.getText().getAsInteger(10, index))
         continue;
       attrs.setOpaqueReturnTypeOf(mangling.str(), index);
+      // SWIFT_ENABLE_TENSORFLOW
+    } else if (attr == TAK_differentiable) {
+      if (arg) {
+        auto argSyntax = arg->getAs<TokenSyntax>();
+        attrs.linear = argSyntax->getTokenKind() == tok::identifier &&
+                       argSyntax->getIdentifierText() == "linear";
+      } else {
+        attrs.linear = false;
+      }
     }
 
     attrs.setAttr(attr, atLoc);

lib/Parse/ParseDecl.cpp

@@ -2536,49 +2536,64 @@ ParsedSyntaxResult<ParsedAttributeSyntax> Parser::parseTypeAttributeSyntax() {
   }
 
   // SWIFT_ENABLE_TENSORFLOW
-  case TAK_differentiable: {
-    bool linear = false;
-    // Check if there is a 'linear' argument.
-    if (Tok.is(tok::l_paren) && peekToken().is(tok::identifier)) {
+  case TAK_differentiable:
+    status |= [&]() -> ParserStatus {
+      // Check if there is a 'linear' argument.
+      // If next tokens are not `'(' identifier`, break early.
+      if (!Tok.is(tok::l_paren) || !peekToken().is(tok::identifier))
+        return makeParserSuccess();
+
       Parser::BacktrackingScope backtrack(*this);
-      consumeToken(tok::l_paren);
+      SourceLoc lParenLoc = Tok.getLoc();
+      auto lParen = consumeTokenSyntax(tok::l_paren);
 
       // Determine if we have '@differentiable(linear) (T) -> U'
       // or '@differentiable (linear) -> U'.
-      if (Tok.getText() == "linear" && consumeIf(tok::identifier)) {
+      if (Tok.getText() == "linear") {
+        auto linearIdentifier = consumeTokenSyntax(tok::identifier);
         if (Tok.is(tok::r_paren) && peekToken().is(tok::l_paren)) {
           // It is being used as an attribute argument, so cancel backtrack
           // as function is linear differentiable.
-          linear = true;
           backtrack.cancelBacktrack();
-          consumeToken(tok::r_paren);
+          builder.useLeftParen(std::move(lParen));
+          builder.useArgument(std::move(linearIdentifier));
+          SourceLoc rParenLoc;
+          auto rParen = parseMatchingTokenSyntax(
+              tok::r_paren, rParenLoc, diag::differentiable_attribute_expected_rparen,
+              lParenLoc);
+          if (!rParen)
+            return makeParserError();
+          builder.useRightParen(std::move(*rParen));
         } else if (Tok.is(tok::l_paren)) {
           // Handle invalid '@differentiable(linear (T) -> U'
           diagnose(Tok, diag::differentiable_attribute_expected_rparen);
           backtrack.cancelBacktrack();
-          status.setIsParseError();
-          break;
+          builder.useLeftParen(std::move(lParen));
+          builder.useArgument(std::move(linearIdentifier));
+          return makeParserError();
         }
       } else if (Tok.is(tok::identifier)) {
         // No 'linear' arg or param type, but now checking if the token is being
         // passed in as an invalid argument to '@differentiable'.
         auto possibleArg = Tok.getText();
         auto t = Tok; // get ref to the argument for clearer diagnostics.
-        consumeToken(tok::identifier);
+        auto argIdentifier = consumeTokenSyntax(tok::identifier);
         // Check if there is an invalid argument getting passed into
         // '@differentiable'.
         if (Tok.is(tok::r_paren) && peekToken().is(tok::l_paren)) {
           // Handling '@differentiable(wrong) (...'.
           diagnose(t, diag::unexpected_argument_differentiable, possibleArg);
-          consumeToken(tok::r_paren);
+          auto rParen = consumeTokenSyntax(tok::r_paren);
           backtrack.cancelBacktrack();
-          status.setIsParseError();
-          break;
+          builder.useLeftParen(std::move(lParen));
+          builder.useArgument(std::move(argIdentifier));
+          builder.useRightParen(std::move(rParen));
+          return makeParserError();
         }
       }
-    }
+      return makeParserSuccess();
+    }();
     break;
-  }
 
   case TAK_convention:
     status |= [&]() -> ParserStatus {