Pull in the upstream @differentiable attribute changes (#28198)

Author: bgogul

include/swift/AST/AutoDiff.h

@@ -141,23 +141,23 @@ class ParsedAutoDiffParameter {
   enum class Kind { Named, Ordered, Self };
 
 private:
-  SourceLoc Loc;
-  Kind Kind;
+  SourceLoc loc;
+  Kind kind;
   union Value {
-    struct { Identifier Name; } Named;
-    struct { unsigned Index; } Ordered;
-    struct {} Self;
+    struct { Identifier name; } Named;
+    struct { unsigned index; } Ordered;
+    struct {} self;
     Value(Identifier name) : Named({name}) {}
     Value(unsigned index) : Ordered({index}) {}
     Value() {}
-  } V;
+  } value;
 
 public:
-  ParsedAutoDiffParameter(SourceLoc loc, enum Kind kind, Value value)
-      : Loc(loc), Kind(kind), V(value) {}
-  
-  ParsedAutoDiffParameter(SourceLoc loc, enum Kind kind, unsigned index)
-      : Loc(loc), Kind(kind), V(index) {}
+  ParsedAutoDiffParameter(SourceLoc loc, Kind kind, Value value)
+    : loc(loc), kind(kind), value(value) {}
+
+  ParsedAutoDiffParameter(SourceLoc loc, Kind kind, unsigned index)
+    : loc(loc), kind(kind), value(index) {}
 
   static ParsedAutoDiffParameter getNamedParameter(SourceLoc loc,
                                                    Identifier name) {
@@ -174,20 +174,20 @@ class ParsedAutoDiffParameter {
   }
 
   Identifier getName() const {
-    assert(Kind == Kind::Named);
-    return V.Named.Name;
+    assert(kind == Kind::Named);
+    return value.Named.name;
   }
 
   unsigned getIndex() const {
-    return V.Ordered.Index;
+    return value.Ordered.index;
   }
 
-  enum Kind getKind() const {
-    return Kind;
+  Kind getKind() const {
+    return kind;
   }
 
   SourceLoc getLoc() const {
-    return Loc;
+    return loc;
   }
 
   bool isEqual(const ParsedAutoDiffParameter &other) const {

include/swift/Parse/Parser.h

@@ -711,6 +711,20 @@ class Parser {
   /// Check whether the current token starts with '>'.
   bool startsWithGreater(Token Tok) { return startsWithSymbol(Tok, '>'); }
 
+  /// Returns true if token is an identifier with the given value.
+  bool isIdentifier(Token Tok, StringRef value) {
+    return Tok.is(tok::identifier) && Tok.getText() == value;
+  }
+
+  /// Returns true if token is the identifier "wrt".
+  bool isWRTIdentifier(Token tok) { return isIdentifier(Tok, "wrt"); }
+
+  /// Returns true if token is the identifier "jvp".
+  bool isJVPIdentifier(Token Tok) { return isIdentifier(Tok, "jvp"); }
+
+  /// Returns true if token is the identifier "vjp".
+  bool isVJPIdentifier(Token Tok) { return isIdentifier(Tok, "vjp"); }
+
   /// Consume the starting '<' of the current token, which may either
   /// be a complete '<' token or some kind of operator token starting with '<',
   /// e.g., '<>'.

lib/Parse/ParseDecl.cpp

@@ -1026,9 +1026,8 @@ bool Parser::parseDifferentiableAttributeArguments(
       return true;
     }
     // Check that token after comma is 'wrt:' or a function specifier label.
-    if (!Tok.is(tok::identifier) || !(Tok.getText() == "wrt" ||
-                                      Tok.getText() == "jvp" ||
-                                      Tok.getText() == "vjp")) {
+    if (!(isWRTIdentifier(Tok) || isJVPIdentifier(Tok) ||
+          isVJPIdentifier(Tok))) {
       diagnose(Tok, diag::attr_differentiable_expected_label);
       return true;
     }
@@ -1047,7 +1046,7 @@ bool Parser::parseDifferentiableAttributeArguments(
     linear = true;
     consumeToken(tok::identifier);
     // If no trailing comma or 'where' clause, terminate parsing arguments.
-    if (Tok.isNot(tok::comma) && Tok.isNot(tok::kw_where))
+    if (Tok.isNot(tok::comma, tok::kw_where))
       return false;
     if (consumeIfTrailingComma())
       return errorAndSkipToEnd();
@@ -1061,11 +1060,11 @@ bool Parser::parseDifferentiableAttributeArguments(
         .fixItReplace(withRespectToRange, "wrt:");
     return errorAndSkipToEnd();
   }
-  if (Tok.is(tok::identifier) && Tok.getText() == "wrt") {
+  if (isWRTIdentifier(Tok)) {
     if (parseDifferentiationParametersClause(params, AttrName))
       return true;
     // If no trailing comma or 'where' clause, terminate parsing arguments.
-    if (Tok.isNot(tok::comma) && Tok.isNot(tok::kw_where))
+    if (Tok.isNot(tok::comma, tok::kw_where))
       return false;
     if (consumeIfTrailingComma())
       return errorAndSkipToEnd();
@@ -1090,7 +1089,7 @@ bool Parser::parseDifferentiableAttributeArguments(
                                  funcDiag, /*allowOperators=*/true,
                                  /*allowZeroArgCompoundNames=*/true);
     // If no trailing comma or 'where' clause, terminate parsing arguments.
-    if (Tok.isNot(tok::comma) && Tok.isNot(tok::kw_where))
+    if (Tok.isNot(tok::comma, tok::kw_where))
       terminateParsingArgs = true;
     return !result.Name;
   };
@@ -1099,7 +1098,7 @@ bool Parser::parseDifferentiableAttributeArguments(
   bool terminateParsingArgs = false;
 
   // Parse 'jvp: <func_name>' (optional).
-  if (Tok.is(tok::identifier) && Tok.getText() == "jvp") {
+  if (isJVPIdentifier(Tok)) {
     SyntaxParsingContext JvpContext(
         SyntaxContext, SyntaxKind::DifferentiableAttributeFuncSpecifier);
     jvpSpec = DeclNameWithLoc();
@@ -1112,7 +1111,7 @@ bool Parser::parseDifferentiableAttributeArguments(
   }
 
   // Parse 'vjp: <func_name>' (optional).
-  if (Tok.is(tok::identifier) && Tok.getText() == "vjp") {
+  if (isVJPIdentifier(Tok)) {
     SyntaxParsingContext VjpContext(
         SyntaxContext, SyntaxKind::DifferentiableAttributeFuncSpecifier);
     vjpSpec = DeclNameWithLoc();
@@ -1127,8 +1126,7 @@ bool Parser::parseDifferentiableAttributeArguments(
   }
 
   // If parser has not advanced and token is not 'where' or ')', emit error.
-  if (Tok.getLoc() == startingLoc &&
-      Tok.isNot(tok::kw_where) && Tok.isNot(tok::r_paren)) {
+  if (Tok.getLoc() == startingLoc && Tok.isNot(tok::kw_where, tok::r_paren)) {
     diagnose(Tok, diag::attr_differentiable_expected_label);
     return errorAndSkipToEnd();
   }

lib/ParseSIL/ParseSIL.cpp

@@ -7023,7 +7023,7 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
     SourceLoc lBraceLoc;
     P.consumeIf(tok::l_brace, lBraceLoc);
     // Parse JVP (optional).
-    if (P.Tok.is(tok::identifier) && P.Tok.getText() == "jvp") {
+    if (P.isJVPIdentifier(P.Tok)) {
       P.consumeToken(tok::identifier);
       if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_token, ":"))
         return true;
@@ -7032,7 +7032,7 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
         return true;
     }
     // Parse VJP (optional).
-    if (P.Tok.is(tok::identifier) && P.Tok.getText() == "vjp") {
+    if (P.isVJPIdentifier(P.Tok)) {
       P.consumeToken(tok::identifier);
       if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_token, ":"))
         return true;