[Autodiff] Derivative Registration for the Get and Set Accessors (#32614)

* initial changes

* Add tests, undo unnecessary changes.

* Fixing up computed properties accessors and adding tests for getters.

* Adding nested type testcase

* Fixing error message for when accessor is referenced but not acutally found.

* Cleanup.

- Improve diagnostic message.
- Clean up code and tests.
- Delete unrelated nested type `@derivative` attribute tests.

* Temporarily disable class subscript setter derivative registration test.

Blocked by SR-13096.

* Adding libsyntax integration and fixing up an error message.

* Added a helper function for checking if the next token is an accessor label.

* Update utils/gyb_syntax_support/AttributeNodes.py

Co-authored-by: Dan Zheng <[email protected]>

* Update lib/Parse/ParseDecl.cpp

Co-authored-by: Dan Zheng <[email protected]>

* Add end-to-end derivative registration tests.

* NFC: run `git clang-format`.

* NFC: clean up formatting.

Re-apply `git clang-format`.

* Clarify parsing ambiguity FIXME comments.

* Adding couple of more testcases and fixing up error message for when accessor is not found on functions resolved.

* Update lib/Sema/TypeCheckAttr.cpp

Co-authored-by: Dan Zheng <[email protected]>

Co-authored-by: Dan Zheng <[email protected]>

Author: ematejska

include/swift/AST/Attr.h

@@ -35,6 +35,7 @@
 #include "swift/AST/Ownership.h"
 #include "swift/AST/PlatformKind.h"
 #include "swift/AST/Requirement.h"
+#include "swift/AST/StorageImpl.h"
 #include "swift/AST/TrailingCallArguments.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -1718,6 +1719,7 @@ class OriginallyDefinedInAttr: public DeclAttribute {
 struct DeclNameRefWithLoc {
   DeclNameRef Name;
   DeclNameLoc Loc;
+  Optional<AccessorKind> AccessorKind;
 };
 
 /// Attribute that marks a function as differentiable.

include/swift/AST/DiagnosticsSema.def

@@ -3095,6 +3095,10 @@ ERROR(derivative_attr_class_member_dynamic_self_result_unsupported,none,
 ERROR(derivative_attr_nonfinal_class_init_unsupported,none,
       "cannot register derivative for 'init' in a non-final class; consider "
       "making %0 final", (Type))
+// TODO(SR-13096): Remove this temporary diagnostic.
+ERROR(derivative_attr_class_setter_unsupported,none,
+      "cannot yet register derivative for class property or subscript setters",
+      ())
 ERROR(derivative_attr_original_already_has_derivative,none,
       "a derivative already exists for %0", (DeclName))
 NOTE(derivative_attr_duplicate_note,none,
@@ -3133,6 +3137,8 @@ NOTE(transpose_attr_wrt_self_self_type_mismatch_note,none,
 ERROR(autodiff_attr_original_decl_invalid_kind,none,
       "%0 is not a 'func', 'init', 'subscript', or 'var' computed property "
       "declaration", (DeclNameRef))
+ERROR(autodiff_attr_accessor_not_found,none,
+      "%0 does not have a '%1' accessor", (DeclNameRef, StringRef))
 ERROR(autodiff_attr_original_decl_none_valid_found,none,
       "could not find function %0 with expected type %1", (DeclNameRef, Type))
 ERROR(autodiff_attr_original_decl_not_same_type_context,none,

lib/Parse/ParseDecl.cpp

@@ -1037,11 +1037,24 @@ bool Parser::parseDifferentiableAttributeArguments(
   return false;
 }
 
+// Helper function that returns the accessor kind if a token is an accessor
+// label.
+static Optional<AccessorKind> isAccessorLabel(const Token &token) {
+  if (token.is(tok::identifier)) {
+    StringRef tokText = token.getText();
+    for (auto accessor : allAccessorKinds())
+      if (tokText == getAccessorLabel(accessor))
+        return accessor;
+  }
+  return None;
+}
+
 /// Helper function that parses 'type-identifier' for `parseQualifiedDeclName`.
 /// Returns true on error. Sets `baseType` to the parsed base type if present,
 /// or to `nullptr` if not. A missing base type is not considered an error.
 static bool parseBaseTypeForQualifiedDeclName(Parser &P, TypeRepr *&baseType) {
   baseType = nullptr;
+  Parser::BacktrackingScope backtrack(P);
 
   // If base type cannot be parsed, return false (no error).
   if (!P.canParseBaseTypeForQualifiedDeclName())
@@ -1057,6 +1070,28 @@ static bool parseBaseTypeForQualifiedDeclName(Parser &P, TypeRepr *&baseType) {
   // `parseTypeIdentifier(/*isParsingQualifiedDeclName*/ true)` leaves the
   // leading period unparsed to avoid syntax verification errors.
   assert(P.startsWithSymbol(P.Tok, '.') && "false");
+
+  // Check if this is a reference to a property or subscript accessor.
+  //
+  // Note: There is an parsing ambiguity here. An accessor label identifier
+  // (e.g. "set") may refer to the final declaration name component instead of
+  // an accessor kind.
+  //
+  // FIXME: It is wrong to backtrack parsing the entire base type if an accessor
+  // label is found. Instead, only the final component of the base type should
+  // be backtracked. It may be best to implement this in
+  // `Parser::parseTypeIdentifier`.
+  //
+  // Example: consider parsing `A.B.property.set`.
+  // Current behavior: base type is entirely backtracked.
+  // Ideal behavior: base type is parsed as `A.B`.
+  if (P.Tok.is(tok::period)) {
+    const Token &nextToken = P.peekToken();
+    if (isAccessorLabel(nextToken).hasValue())
+      return false;
+  }
+
+  backtrack.cancelBacktrack();
   P.consumeStartingCharacterOfCurrentToken(tok::period);
 
   // Set base type and return false (no error).
@@ -1079,20 +1114,52 @@ static bool parseBaseTypeForQualifiedDeclName(Parser &P, TypeRepr *&baseType) {
 static bool parseQualifiedDeclName(Parser &P, Diag<> nameParseError,
                                    TypeRepr *&baseType,
                                    DeclNameRefWithLoc &original) {
-  SyntaxParsingContext DeclNameContext(P.SyntaxContext,
-                                       SyntaxKind::QualifiedDeclName);
-  // Parse base type.
-  if (parseBaseTypeForQualifiedDeclName(P, baseType))
-    return true;
-  // Parse final declaration name.
-  original.Name = P.parseDeclNameRef(
-      original.Loc, nameParseError,
-      Parser::DeclNameFlag::AllowZeroArgCompoundNames |
-      Parser::DeclNameFlag::AllowKeywordsUsingSpecialNames |
-      Parser::DeclNameFlag::AllowOperators);
-  // The base type is optional, but the final unqualified declaration name is
-  // not. If name could not be parsed, return true for error.
-  return !original.Name;
+  {
+    SyntaxParsingContext DeclNameContext(P.SyntaxContext,
+                                         SyntaxKind::QualifiedDeclName);
+    // Parse base type.
+    if (parseBaseTypeForQualifiedDeclName(P, baseType))
+      return true;
+    // Parse final declaration name.
+    original.Name = P.parseDeclNameRef(
+        original.Loc, nameParseError,
+        Parser::DeclNameFlag::AllowZeroArgCompoundNames |
+            Parser::DeclNameFlag::AllowKeywordsUsingSpecialNames |
+            Parser::DeclNameFlag::AllowOperators);
+    // The base type is optional, but the final unqualified declaration name is
+    // not. If name could not be parsed, return true for error.
+    if (!original.Name)
+      return true;
+  }
+
+  // Parse an optional accessor kind.
+  //
+  // Note: there is an parsing ambiguity here.
+  //
+  // Example: `A.B.property.set` may be parsed as one of the following:
+  //
+  // 1. No accessor kind.
+  // - Base type: `A.B.property`
+  // - Declaration name: `set`
+  // - Accessor kind: <none>
+  //
+  // 2. Accessor kind exists.
+  // - Base type: `A.B`
+  // - Declaration name: `property`
+  // - Accessor kind: `set`
+  //
+  // Currently, we follow (2) because it's more useful in practice.
+  if (P.Tok.is(tok::period)) {
+    const Token &nextToken = P.peekToken();
+    Optional<AccessorKind> kind = isAccessorLabel(nextToken);
+    if (kind.hasValue()) {
+      original.AccessorKind = kind;
+      P.consumeIf(tok::period);
+      P.consumeIf(tok::identifier);
+    }
+  }
+
+  return false;
 }
 
 /// Parse a `@derivative(of:)` attribute, returning true on error.

lib/Sema/TypeCheckAttr.cpp

@@ -30,6 +30,7 @@
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/StorageImpl.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 #include "swift/Parse/Lexer.h"
@@ -3609,12 +3610,14 @@ static IndexSubset *computeDifferentiabilityParameters(
 // If the function declaration cannot be resolved, emits a diagnostic and
 // returns nullptr.
 static AbstractFunctionDecl *findAbstractFunctionDecl(
-    DeclNameRef funcName, SourceLoc funcNameLoc, Type baseType,
+    DeclNameRef funcName, SourceLoc funcNameLoc,
+    Optional<AccessorKind> accessorKind, Type baseType,
     DeclContext *lookupContext,
     const std::function<bool(AbstractFunctionDecl *)> &isValidCandidate,
     const std::function<void()> &noneValidDiagnostic,
     const std::function<void()> &ambiguousDiagnostic,
     const std::function<void()> &notFunctionDiagnostic,
+    const std::function<void()> &missingAccessorDiagnostic,
     NameLookupOptions lookupOptions,
     const Optional<std::function<bool(AbstractFunctionDecl *)>>
         &hasValidTypeCtx,
@@ -3640,6 +3643,7 @@ static AbstractFunctionDecl *findAbstractFunctionDecl(
   bool wrongTypeContext = false;
   bool ambiguousFuncDecl = false;
   bool foundInvalid = false;
+  bool missingAccessor = false;
 
   // Filter lookup results.
   for (auto choice : results) {
@@ -3648,10 +3652,21 @@ static AbstractFunctionDecl *findAbstractFunctionDecl(
       continue;
     // Cast the candidate to an `AbstractFunctionDecl`.
     auto *candidate = dyn_cast<AbstractFunctionDecl>(decl);
-    // If the candidate is an `AbstractStorageDecl`, use its getter as the
-    // candidate.
-    if (auto *asd = dyn_cast<AbstractStorageDecl>(decl))
-      candidate = asd->getOpaqueAccessor(AccessorKind::Get);
+    // If the candidate is an `AbstractStorageDecl`, use one of its accessors as
+    // the candidate.
+    if (auto *asd = dyn_cast<AbstractStorageDecl>(decl)) {
+      // If accessor kind is specified, use corresponding accessor from the
+      // candidate. Otherwise, use the getter by default.
+      if (accessorKind != None) {
+        candidate = asd->getOpaqueAccessor(accessorKind.getValue());
+        // Error if candidate is missing the requested accessor.
+        if (!candidate)
+          missingAccessor = true;
+      } else
+        candidate = asd->getOpaqueAccessor(AccessorKind::Get);
+    } else if (accessorKind != None) {
+      missingAccessor = true;
+    }
     if (!candidate) {
       notFunction = true;
       continue;
@@ -3671,8 +3686,9 @@ static AbstractFunctionDecl *findAbstractFunctionDecl(
     }
     resolvedCandidate = candidate;
   }
+
   // If function declaration was resolved, return it.
-  if (resolvedCandidate)
+  if (resolvedCandidate && !missingAccessor)
     return resolvedCandidate;
 
   // Otherwise, emit the appropriate diagnostic and return nullptr.
@@ -3685,6 +3701,10 @@ static AbstractFunctionDecl *findAbstractFunctionDecl(
     ambiguousDiagnostic();
     return nullptr;
   }
+  if (missingAccessor) {
+    missingAccessorDiagnostic();
+    return nullptr;
+  }
   if (wrongTypeContext) {
     assert(invalidTypeCtxDiagnostic &&
            "Type context diagnostic should've been specified");
@@ -4429,6 +4449,13 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
                    diag::autodiff_attr_original_decl_invalid_kind,
                    originalName.Name);
   };
+  auto missingAccessorDiagnostic = [&]() {
+    auto accessorKind = originalName.AccessorKind.getValueOr(AccessorKind::Get);
+    auto accessorLabel = getAccessorLabel(accessorKind);
+    diags.diagnose(originalName.Loc, diag::autodiff_attr_accessor_not_found,
+                   originalName.Name, accessorLabel);
+  };
+
   std::function<void()> invalidTypeContextDiagnostic = [&]() {
     diags.diagnose(originalName.Loc,
                    diag::autodiff_attr_original_decl_not_same_type_context,
@@ -4473,15 +4500,17 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
 
   // Look up original function.
   auto *originalAFD = findAbstractFunctionDecl(
-      originalName.Name, originalName.Loc.getBaseNameLoc(), baseType,
-      derivativeTypeCtx, isValidOriginal, noneValidDiagnostic,
-      ambiguousDiagnostic, notFunctionDiagnostic, lookupOptions,
-      hasValidTypeContext, invalidTypeContextDiagnostic);
+      originalName.Name, originalName.Loc.getBaseNameLoc(),
+      originalName.AccessorKind, baseType, derivativeTypeCtx, isValidOriginal,
+      noneValidDiagnostic, ambiguousDiagnostic, notFunctionDiagnostic,
+      missingAccessorDiagnostic, lookupOptions, hasValidTypeContext,
+      invalidTypeContextDiagnostic);
   if (!originalAFD)
     return true;
-  // Diagnose original stored properties. Stored properties cannot have custom
-  // registered derivatives.
+
   if (auto *accessorDecl = dyn_cast<AccessorDecl>(originalAFD)) {
+    // Diagnose original stored properties. Stored properties cannot have custom
+    // registered derivatives.
     auto *asd = accessorDecl->getStorage();
     if (asd->hasStorage()) {
       diags.diagnose(originalName.Loc,
@@ -4491,6 +4520,17 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
                      asd->getName());
       return true;
     }
+    // Diagnose original class property and subscript setters.
+    // TODO(SR-13096): Fix derivative function typing results regarding
+    // class-typed function parameters.
+    if (asd->getDeclContext()->getSelfClassDecl() &&
+        accessorDecl->getAccessorKind() == AccessorKind::Set) {
+      diags.diagnose(originalName.Loc,
+                     diag::derivative_attr_class_setter_unsupported);
+      diags.diagnose(originalAFD->getLoc(), diag::decl_declared_here,
+                     asd->getName());
+      return true;
+    }
   }
   // Diagnose if original function is an invalid class member.
   bool isOriginalClassMember =
@@ -4998,6 +5038,13 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
              diag::autodiff_attr_original_decl_invalid_kind,
              originalName.Name);
   };
+  auto missingAccessorDiagnostic = [&]() {
+    auto accessorKind = originalName.AccessorKind.getValueOr(AccessorKind::Get);
+    auto accessorLabel = getAccessorLabel(accessorKind);
+    diagnose(originalName.Loc, diag::autodiff_attr_accessor_not_found,
+             originalName.Name, accessorLabel);
+  };
+
   std::function<void()> invalidTypeContextDiagnostic = [&]() {
     diagnose(originalName.Loc,
              diag::autodiff_attr_original_decl_not_same_type_context,
@@ -5028,8 +5075,9 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
   if (attr->getBaseTypeRepr())
     funcLoc = attr->getBaseTypeRepr()->getLoc();
   auto *originalAFD = findAbstractFunctionDecl(
-      originalName.Name, funcLoc, baseType, transposeTypeCtx, isValidOriginal,
-      noneValidDiagnostic, ambiguousDiagnostic, notFunctionDiagnostic,
+      originalName.Name, funcLoc, originalName.AccessorKind, baseType,
+      transposeTypeCtx, isValidOriginal, noneValidDiagnostic,
+      ambiguousDiagnostic, notFunctionDiagnostic, missingAccessorDiagnostic,
       lookupOptions, hasValidTypeContext, invalidTypeContextDiagnostic);
   if (!originalAFD) {
     attr->setInvalid();

lib/Serialization/Deserialization.cpp

@@ -4388,8 +4388,8 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
             scratch, isImplicit, origNameId, origDeclId, rawDerivativeKind,
             parameters);
 
-        DeclNameRefWithLoc origName{
-            DeclNameRef(MF.getDeclBaseName(origNameId)), DeclNameLoc()};
+        DeclNameRefWithLoc origName{DeclNameRef(MF.getDeclBaseName(origNameId)),
+                                    DeclNameLoc(), None};
         auto derivativeKind =
             getActualAutoDiffDerivativeFunctionKind(rawDerivativeKind);
         if (!derivativeKind)
@@ -4418,7 +4418,7 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
             scratch, isImplicit, origNameId, origDeclId, parameters);
 
         DeclNameRefWithLoc origName{
-            DeclNameRef(MF.getDeclBaseName(origNameId)), DeclNameLoc()};
+            DeclNameRef(MF.getDeclBaseName(origNameId)), DeclNameLoc(), None};
         auto *origDecl = cast<AbstractFunctionDecl>(MF.getDecl(origDeclId));
         llvm::SmallBitVector parametersBitVector(parameters.size());
         for (unsigned i : indices(parameters))

test/AutoDiff/Parse/derivative_attr_parse.swift

@@ -26,6 +26,27 @@ func dfoo(x: Float) -> (value: Float, differential: (Float) -> (Float)) {
   return (x, { $0 })
 }
 
+@derivative(of: property.get) // ok
+func dPropertyGetter() -> ()
+
+@derivative(of: subscript.get) // ok
+func dSubscriptGetter() -> ()
+
+@derivative(of: subscript(_:label:).get) // ok
+func dLabeledSubscriptGetter() -> ()
+
+@derivative(of: property.set) // ok
+func dPropertySetter() -> ()
+
+@derivative(of: subscript.set) // ok
+func dSubscriptSetter() -> ()
+
+@derivative(of: subscript(_:label:).set) // ok
+func dLabeledSubscriptSetter() -> ()
+
+@derivative(of: nestedType.name) // ok
+func dNestedTypeFunc() -> ()
+
 /// Bad
 
 // expected-error @+2 {{expected an original function name}}
@@ -98,3 +119,16 @@ func testLocalDerivativeRegistration() {
   @derivative(of: sin)
   func dsin()
 }
+
+
+func testLocalDerivativeRegistration() {
+  // expected-error @+1 {{attribute '@derivative' can only be used in a non-local scope}}
+  @derivative(of: sin)
+  func dsin()
+}
+
+// expected-error @+2 {{expected ',' separator}}
+// expected-error @+1 {{expected declaration}}
+@derivative(of: nestedType.name.set)
+func dNestedTypePropertySetter() -> ()
+