[AutoDiff] Support derivative registration for more declaration kinds. (swiftlang#28468)

Make `@differentiating` and `@transposing` attributes support more original
declaration kinds: computed properties, subscripts, and initializers.

- Change `DifferentiatingAttr` and `TransposingAttr` to store an
  `AbstractFunctionDecl` representing the original declaration, instead of a
  `FuncDecl`.
- Change attribute parsing to support initializer/subscript `DeclName`s.
- Make `TypeChecker::lookupFuncDecl` a static function in TypeCheckAttr.cpp.
- Assorted parsing and type-checking gardening.

`@differentiating` now has feature parity with
`@differentiable(jvp: ..., vjp: ...)` for derivative registration.

This is a necessary step towards making `@differentiating` and `@transposing`
the canonical mechanism for registering derivative/transpose functions.

Registering non-`func` declaration derivatives with `@differentiable` attribute
`jvp:`/`vjp:` labels is now explicitly rejected.

Resolves TF-281.

Todos:
- TF-997: support `@transposing` attribute with initializer original
  declarations.
- TF-988: do not reuse `@differentiable` attribute type-checking diagnostics
  for `@differentiating`/`@transposing` attribute type-checking.

Author: dan-zheng

include/swift/AST/Attr.h

@@ -1679,9 +1679,9 @@ class DifferentiatingAttr final
   friend TrailingObjects;
 
   /// The original function name.
-  DeclNameWithLoc Original;
-  /// The original function, resolved by the type checker.
-  FuncDecl *OriginalFunction = nullptr;
+  DeclNameWithLoc OriginalFunctionName;
+  /// The original function declaration, resolved by the type checker.
+  AbstractFunctionDecl *OriginalFunction = nullptr;
   /// The number of parsed parameters specified in 'wrt:'.
   unsigned NumParsedParameters = 0;
   /// The differentiation parameters' indices, resolved by the type checker.
@@ -1706,9 +1706,15 @@ class DifferentiatingAttr final
                                      DeclNameWithLoc original,
                                      IndexSubset *indices);
 
-  DeclNameWithLoc getOriginal() const { return Original; }
-  FuncDecl *getOriginalFunction() const { return OriginalFunction; }
-  void setOriginalFunction(FuncDecl *decl) { OriginalFunction = decl; }
+  DeclNameWithLoc getOriginalFunctionName() const {
+    return OriginalFunctionName;
+  }
+  AbstractFunctionDecl *getOriginalFunction() const {
+    return OriginalFunction;
+  }
+  void setOriginalFunction(AbstractFunctionDecl *decl) {
+    OriginalFunction = decl;
+  }
 
   /// The parsed differentiation parameters, i.e. the list of parameters
   /// specified in 'wrt:'.
@@ -1750,9 +1756,9 @@ class TransposingAttr final
   /// is an instance/static method).
   TypeRepr *BaseType;
   /// The original function name.
-  DeclNameWithLoc Original;
-  /// The original function, resolved by the type checker.
-  FuncDecl *OriginalFunction = nullptr;
+  DeclNameWithLoc OriginalFunctionName;
+  /// The original function declaration, resolved by the type checker.
+  AbstractFunctionDecl *OriginalFunction = nullptr;
   /// The number of parsed parameters specified in 'wrt:'.
   unsigned NumParsedParameters = 0;
   /// The differentiation parameters' indices, resolved by the type checker.
@@ -1779,10 +1785,15 @@ class TransposingAttr final
                                  IndexSubset *indices);
 
   TypeRepr *getBaseType() const { return BaseType; }
-  DeclNameWithLoc getOriginal() const { return Original; }
-
-  FuncDecl *getOriginalFunction() const { return OriginalFunction; }
-  void setOriginalFunction(FuncDecl *decl) { OriginalFunction = decl; }
+  DeclNameWithLoc getOriginalFunctionName() const {
+    return OriginalFunctionName;
+  }
+  AbstractFunctionDecl *getOriginalFunction() const {
+    return OriginalFunction;
+  }
+  void setOriginalFunction(AbstractFunctionDecl *decl) {
+    OriginalFunction = decl;
+  }
 
   /// The parsed transposing parameters, i.e. the list of parameters
   /// specified in 'wrt:'.

include/swift/AST/DiagnosticsSema.def

@@ -2954,9 +2954,8 @@ NOTE(differentiable_attr_duplicate_note,none,
      "other attribute declared here", ())
 ERROR(differentiable_attr_function_not_same_type_context,none,
       "%0 is not defined in the current type context", (DeclName))
-ERROR(differentiable_attr_specified_not_function,none,
-      "%0 is not a function to be used as derivative function",
-      (DeclName))
+ERROR(differentiable_attr_derivative_not_function,none,
+      "registered derivative %0 must be a 'func' declaration", (DeclName))
 ERROR(differentiable_attr_class_derivative_not_final,none,
       "class member derivative must be final", ())
 ERROR(differentiable_attr_ambiguous_function_identifier,none,
@@ -3020,6 +3019,8 @@ ERROR(differentiating_attr_overload_not_found,none,
       "could not find function %0 with expected type %1", (DeclName, Type))
 ERROR(differentiating_attr_not_in_same_file_as_original,none,
       "derivative not in the same file as the original function", ())
+ERROR(differentiating_attr_original_stored_property_unsupported,none,
+      "cannot register derivative for stored property %0", (DeclName))
 ERROR(differentiating_attr_original_already_has_derivative,none,
       "a derivative already exists for %0", (DeclName))
 
@@ -3033,7 +3034,7 @@ ERROR(transposing_attr_cannot_use_named_wrt_params,none,
       "cannot use named 'wrt' parameters in '@transposing' attribute, found %0",
       (Identifier))
 ERROR(transposing_attr_result_value_not_differentiable,none,
-      "'@transposing' attribute requires original function result to "
+      "'@transposing' attribute requires original function result %0 to "
       "conform to 'Differentiable'", (Type))
 
 // differentiation `wrt` parameters clause

lib/AST/Attr.cpp

@@ -926,7 +926,7 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer << '(';
     auto *attr = cast<DifferentiatingAttr>(this);
     auto *derivative = cast<AbstractFunctionDecl>(D);
-    Printer << attr->getOriginal().Name;
+    Printer << attr->getOriginalFunctionName().Name;
     auto diffParamsString = getDifferentiationParametersClauseString(
         derivative, attr->getParameterIndices(), attr->getParsedParameters());
     if (!diffParamsString.empty())
@@ -941,7 +941,7 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer << '(';
     auto *attr = cast<TransposingAttr>(this);
     auto *transpose = cast<AbstractFunctionDecl>(D);
-    Printer << attr->getOriginal().Name;
+    Printer << attr->getOriginalFunctionName().Name;
     auto transParamsString = getTransposedParametersClauseString(
         transpose, attr->getParameterIndices(), attr->getParsedParameters());
     if (!transParamsString.empty())
@@ -1570,19 +1570,21 @@ void DifferentiableAttr::print(llvm::raw_ostream &OS, const Decl *D,
 // SWIFT_ENABLE_TENSORFLOW
 DifferentiatingAttr::DifferentiatingAttr(
     bool implicit, SourceLoc atLoc, SourceRange baseRange,
-    DeclNameWithLoc original, ArrayRef<ParsedAutoDiffParameter> params)
+    DeclNameWithLoc originalName, ArrayRef<ParsedAutoDiffParameter> params)
     : DeclAttribute(DAK_Differentiating, atLoc, baseRange, implicit),
-      Original(std::move(original)), NumParsedParameters(params.size()) {
+      OriginalFunctionName(std::move(originalName)),
+      NumParsedParameters(params.size()) {
   std::copy(params.begin(), params.end(),
             getTrailingObjects<ParsedAutoDiffParameter>());
 }
 
 DifferentiatingAttr::DifferentiatingAttr(bool implicit, SourceLoc atLoc,
                                          SourceRange baseRange,
-                                         DeclNameWithLoc original,
+                                         DeclNameWithLoc originalName,
                                          IndexSubset *indices)
     : DeclAttribute(DAK_Differentiating, atLoc, baseRange, implicit),
-      Original(std::move(original)), ParameterIndices(indices) {}
+      OriginalFunctionName(std::move(originalName)), ParameterIndices(indices) {
+}
 
 DifferentiatingAttr *
 DifferentiatingAttr::create(ASTContext &context, bool implicit, SourceLoc atLoc,
@@ -1607,20 +1609,21 @@ DifferentiatingAttr *DifferentiatingAttr::create(ASTContext &context,
 
 TransposingAttr::TransposingAttr(bool implicit, SourceLoc atLoc,
                                  SourceRange baseRange, TypeRepr *baseType,
-                                 DeclNameWithLoc original,
+                                 DeclNameWithLoc originalName,
                                  ArrayRef<ParsedAutoDiffParameter> params)
     : DeclAttribute(DAK_Transposing, atLoc, baseRange, implicit),
-      BaseType(baseType), Original(std::move(original)),
+      BaseType(baseType), OriginalFunctionName(std::move(originalName)),
       NumParsedParameters(params.size()) {
   std::uninitialized_copy(params.begin(), params.end(),
                           getTrailingObjects<ParsedAutoDiffParameter>());
 }
 
 TransposingAttr::TransposingAttr(bool implicit, SourceLoc atLoc,
                                  SourceRange baseRange, TypeRepr *baseType,
-                                 DeclNameWithLoc original, IndexSubset *indices)
+                                 DeclNameWithLoc originalName,
+                                 IndexSubset *indices)
     : DeclAttribute(DAK_Transposing, atLoc, baseRange, implicit),
-      BaseType(baseType), Original(std::move(original)),
+      BaseType(baseType), OriginalFunctionName(std::move(originalName)),
       ParameterIndices(indices) {}
 
 TransposingAttr *

lib/Parse/ParseDecl.cpp

@@ -1069,7 +1069,7 @@ bool Parser::parseDifferentiableAttributeArguments(
     SyntaxParsingContext FuncDeclNameContext(
          SyntaxContext, SyntaxKind::FunctionDeclName);
     Diagnostic funcDiag(diag::attr_differentiable_expected_function_name.ID,
-                        { label });
+                        {label});
     result.Name =
         parseUnqualifiedDeclName(/*afterDot=*/false, result.Loc,
                                  funcDiag, /*allowOperators=*/true,
@@ -1165,11 +1165,14 @@ Parser::parseDifferentiatingAttribute(SourceLoc atLoc, SourceLoc loc) {
       // Parse the name of the function.
       SyntaxParsingContext FuncDeclNameContext(
           SyntaxContext, SyntaxKind::FunctionDeclName);
+      // NOTE: Use `afterDot = true` and `allowDeinitAndSubscript = true` to
+      // enable, e.g. `@differentiating(init)` and
+      // `@differentiating(subscript)`.
       original.Name = parseUnqualifiedDeclName(
-          /*afterDot*/ false, original.Loc,
+          /*afterDot*/ true, original.Loc,
           diag::attr_differentiating_expected_original_name,
-          /*allowOperators*/ true, /*allowZeroArgCompoundNames*/ true);
-
+          /*allowOperators*/ true, /*allowZeroArgCompoundNames*/ true,
+          /*allowDeinitAndSubscript*/ true);
       if (consumeIfTrailingComma())
         return makeParserError();
     }
@@ -1228,19 +1231,13 @@ bool parseQualifiedDeclName(Parser &P, Diag<> nameParseError,
   if (parseBaseTypeForQualifiedDeclName(P, baseType))
     return true;
 
-  // If base type was parsed and has at least one component, then there was a
-  // dot before the current token.
-  bool afterDot = false;
-  if (baseType) {
-    if (auto ident = dyn_cast<IdentTypeRepr>(baseType)) {
-      auto components = ident->getComponentRange();
-      afterDot = std::distance(components.begin(), components.end()) > 0;
-    }
-  }
+  // NOTE: Use `afterDot = true` and `allowDeinitAndSubscript = true` to enable
+  // initializer and subscript lookup.
   original.Name =
-      P.parseUnqualifiedDeclName(afterDot, original.Loc, nameParseError,
-                                 /*allowOperators*/ true,
-                                 /*allowZeroArgCompoundNames*/ true);
+      P.parseUnqualifiedDeclName(/*afterDot*/ true, original.Loc,
+                                 nameParseError, /*allowOperators*/ true,
+                                 /*allowZeroArgCompoundNames*/ true,
+                                 /*allowDeinitAndSubscript*/ true);
 
   // The base type is optional, but the final unqualified decl name is not.
   // If name could not be parsed, return true for error.
@@ -1285,7 +1282,6 @@ ParserResult<TransposingAttr> Parser::parseTransposingAttribute(SourceLoc atLoc,
                                  diag::attr_transposing_expected_original_name,
                                  baseType, original))
         return makeParserError();
-      
       if (consumeIfTrailingComma())
         return makeParserError();
     }