[AutoDiff] factor derivative typechecking helper out of AttributeChecker

include/swift/AST/ASTContext.h

@@ -290,7 +290,7 @@ class ASTContext final {
   // derivative function configurations per original `AbstractFunctionDecl`.
   llvm::DenseMap<
       std::tuple<Decl *, IndexSubset *, AutoDiffDerivativeFunctionKind>,
-      DerivativeAttr *>
+      llvm::SmallPtrSet<DerivativeAttr *, 1>>
       DerivativeAttrs;
 
 private:

lib/Sema/TypeCheckAttr.cpp

@@ -3391,14 +3391,28 @@ getAutoDiffOriginalFunctionType(AnyFunctionType *derivativeFnTy) {
   return originalType;
 }
 
-void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
+/// Typechecks the given derivative attribute `attr` on decl `D`.
+///
+/// Effects are:
+/// - Sets the original function and parameter indices on `attr`.
+/// - Diagnoses errors.
+/// - Stores the attribute in `ASTContext::DerivativeAttrs`.
+///
+/// \returns true on error, false on success.
+static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
+                                    DerivativeAttr *attr) {
+  // Note: Implementation must be idempotent because it can get called multiple
+  // times for the same attribute.
+
+  auto &diags = Ctx.Diags;
+
   // `@derivative` attribute requires experimental differentiable programming
   // to be enabled.
   auto &ctx = D->getASTContext();
   if (!ctx.LangOpts.EnableExperimentalDifferentiableProgramming) {
-    diagnoseAndRemoveAttr(
-        attr, diag::experimental_differentiable_programming_disabled);
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::experimental_differentiable_programming_disabled);
+    return true;
   }
   auto *derivative = cast<FuncDecl>(D);
   auto lookupConformance =
@@ -3418,26 +3432,27 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   auto derivativeResultTupleType = derivativeResultType->getAs<TupleType>();
   if (!derivativeResultTupleType ||
       derivativeResultTupleType->getNumElements() != 2) {
-    diagnoseAndRemoveAttr(attr, diag::derivative_attr_expected_result_tuple);
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_expected_result_tuple);
+    return true;
   }
   auto valueResultElt = derivativeResultTupleType->getElement(0);
   auto funcResultElt = derivativeResultTupleType->getElement(1);
   // Get derivative kind and derivative function identifier.
   AutoDiffDerivativeFunctionKind kind;
   if (valueResultElt.getName().str() != "value") {
-    diagnoseAndRemoveAttr(
-        attr, diag::derivative_attr_invalid_result_tuple_value_label);
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_invalid_result_tuple_value_label);
+    return true;
   }
   if (funcResultElt.getName().str() == "differential") {
     kind = AutoDiffDerivativeFunctionKind::JVP;
   } else if (funcResultElt.getName().str() == "pullback") {
     kind = AutoDiffDerivativeFunctionKind::VJP;
   } else {
-    diagnoseAndRemoveAttr(
-        attr, diag::derivative_attr_invalid_result_tuple_func_label);
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_invalid_result_tuple_func_label);
+    return true;
   }
   attr->setDerivativeKind(kind);
   // `value: R` result tuple element must conform to `Differentiable`.
@@ -3448,10 +3463,10 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   auto valueResultConf = TypeChecker::conformsToProtocol(
       valueResultType, diffableProto, derivative->getDeclContext(), None);
   if (!valueResultConf) {
-    diagnoseAndRemoveAttr(attr,
-                          diag::derivative_attr_result_value_not_differentiable,
-                          valueResultElt.getType());
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_result_value_not_differentiable,
+                   valueResultElt.getType());
+    return true;
   }
 
   // Compute expected original function type and look up original function.
@@ -3495,22 +3510,23 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   };
 
   auto noneValidDiagnostic = [&]() {
-    diagnose(originalName.Loc,
-             diag::autodiff_attr_original_decl_none_valid_found,
-             originalName.Name, originalFnType);
+    diags.diagnose(originalName.Loc,
+                   diag::autodiff_attr_original_decl_none_valid_found,
+                   originalName.Name, originalFnType);
   };
   auto ambiguousDiagnostic = [&]() {
-    diagnose(originalName.Loc, diag::attr_ambiguous_reference_to_decl,
-             originalName.Name, attr->getAttrName());
+    diags.diagnose(originalName.Loc, diag::attr_ambiguous_reference_to_decl,
+                   originalName.Name, attr->getAttrName());
   };
   auto notFunctionDiagnostic = [&]() {
-    diagnose(originalName.Loc, diag::autodiff_attr_original_decl_invalid_kind,
-             originalName.Name);
+    diags.diagnose(originalName.Loc,
+                   diag::autodiff_attr_original_decl_invalid_kind,
+                   originalName.Name);
   };
   std::function<void()> invalidTypeContextDiagnostic = [&]() {
-    diagnose(originalName.Loc,
-             diag::autodiff_attr_original_decl_not_same_type_context,
-             originalName.Name);
+    diags.diagnose(originalName.Loc,
+                   diag::autodiff_attr_original_decl_not_same_type_context,
+                   originalName.Name);
   };
 
   // Returns true if the derivative function and original function candidate are
@@ -3544,22 +3560,19 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
       derivativeTypeCtx, isValidOriginal, noneValidDiagnostic,
       ambiguousDiagnostic, notFunctionDiagnostic, lookupOptions,
       hasValidTypeContext, invalidTypeContextDiagnostic);
-  if (!originalAFD) {
-    attr->setInvalid();
-    return;
-  }
+  if (!originalAFD)
+    return true;
   // Diagnose original stored properties. Stored properties cannot have custom
   // registered derivatives.
   if (auto *accessorDecl = dyn_cast<AccessorDecl>(originalAFD)) {
     auto *asd = accessorDecl->getStorage();
     if (asd->hasStorage()) {
-      diagnose(originalName.Loc,
-               diag::derivative_attr_original_stored_property_unsupported,
-               originalName.Name);
-      diagnose(originalAFD->getLoc(), diag::decl_declared_here,
-               asd->getFullName());
-      attr->setInvalid();
-      return;
+      diags.diagnose(originalName.Loc,
+                     diag::derivative_attr_original_stored_property_unsupported,
+                     originalName.Name);
+      diags.diagnose(originalAFD->getLoc(), diag::decl_declared_here,
+                     asd->getFullName());
+      return true;
     }
   }
   attr->setOriginalFunction(originalAFD);
@@ -3576,19 +3589,15 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
     resolvedWrtParamIndices = computeDifferentiationParameters(
         parsedWrtParams, derivative, derivative->getGenericEnvironment(),
         attr->getAttrName(), attr->getLocation());
-  if (!resolvedWrtParamIndices) {
-    attr->setInvalid();
-    return;
-  }
+  if (!resolvedWrtParamIndices)
+    return true;
 
   // Check if the `wrt:` parameter indices are valid.
   if (checkDifferentiationParameters(
           originalAFD, resolvedWrtParamIndices, originalFnType,
           derivative->getGenericEnvironment(), derivative->getModuleContext(),
-          parsedWrtParams, attr->getLocation())) {
-    attr->setInvalid();
-    return;
-  }
+          parsedWrtParams, attr->getLocation()))
+    return true;
 
   // Set the resolved `wrt:` parameter indices in the attribute.
   attr->setParameterIndices(resolvedWrtParamIndices);
@@ -3647,42 +3656,56 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   // Check if differential/pullback type matches expected type.
   if (!actualFuncEltType->isEqual(expectedFuncEltType)) {
     // Emit differential/pullback type mismatch error on attribute.
-    diagnoseAndRemoveAttr(attr, diag::derivative_attr_result_func_type_mismatch,
-                          funcResultElt.getName(), originalAFD->getFullName());
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_result_func_type_mismatch,
+                   funcResultElt.getName(), originalAFD->getFullName());
     // Emit note with expected differential/pullback type on actual type
     // location.
     auto *tupleReturnTypeRepr =
         cast<TupleTypeRepr>(derivative->getBodyResultTypeLoc().getTypeRepr());
     auto *funcEltTypeRepr = tupleReturnTypeRepr->getElementType(1);
-    diagnose(funcEltTypeRepr->getStartLoc(),
-             diag::derivative_attr_result_func_type_mismatch_note,
-             funcResultElt.getName(), expectedFuncEltType)
+    diags
+        .diagnose(funcEltTypeRepr->getStartLoc(),
+                  diag::derivative_attr_result_func_type_mismatch_note,
+                  funcResultElt.getName(), expectedFuncEltType)
         .highlight(funcEltTypeRepr->getSourceRange());
     // Emit note showing original function location, if possible.
     if (originalAFD->getLoc().isValid())
-      diagnose(originalAFD->getLoc(),
-               diag::derivative_attr_result_func_original_note,
-               originalAFD->getFullName());
-    return;
+      diags.diagnose(originalAFD->getLoc(),
+                     diag::derivative_attr_result_func_original_note,
+                     originalAFD->getFullName());
+    return true;
   }
 
   // Reject different-file derivative registration.
   // TODO(TF-1021): Lift this restriction.
   if (originalAFD->getParentSourceFile() != derivative->getParentSourceFile()) {
-    diagnoseAndRemoveAttr(attr,
-                          diag::derivative_attr_not_in_same_file_as_original);
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_not_in_same_file_as_original);
+    return true;
   }
 
   // Reject duplicate `@derivative` attributes.
-  auto insertion = Ctx.DerivativeAttrs.try_emplace(
-      std::make_tuple(originalAFD, resolvedWrtParamIndices, kind), attr);
-  if (!insertion.second) {
-    diagnoseAndRemoveAttr(attr,
-                          diag::derivative_attr_original_already_has_derivative,
-                          originalAFD->getFullName());
-    diagnose(insertion.first->getSecond()->getLocation(),
-             diag::derivative_attr_duplicate_note);
-    return;
+  auto &derivativeAttrs = Ctx.DerivativeAttrs[std::make_tuple(
+      originalAFD, resolvedWrtParamIndices, kind)];
+  derivativeAttrs.insert(attr);
+  if (derivativeAttrs.size() > 1) {
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_original_already_has_derivative,
+                   originalAFD->getFullName());
+    for (auto *duplicateAttr : derivativeAttrs) {
+      if (duplicateAttr == attr)
+        continue;
+      diags.diagnose(duplicateAttr->getLocation(),
+                     diag::derivative_attr_duplicate_note);
+    }
+    return true;
   }
+
+  return false;
+}
+
+void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
+  if (typeCheckDerivativeAttr(Ctx, D, attr))
+    attr->setInvalid();
 }