[Sema] [AutoDiff] Modify/fix Differentiable derived conformances. (#21871)

- Fix crasher regarding bad interaction between AdditiveArithmetic and
  Differentiable derived conformances, where both code paths attempt to
  synthesize memberwise initializers.

- Conform synthesized `AllDifferentiableVariables` struct to `KeyPathIterable`
  if parent conforms to `KeyPathIterable`.
  - This is important for key-path based optimizer definitions.

- Modify associated struct synthesis to check whether any `@noDerivative`
  stored properties exist.
  - If any `@noDerivative` stored properties exist, `AllDifferentiableVariables`
    struct must be synthesized.

- Refactor derived conformances code, add tests.

Author: dan-zheng

include/swift/AST/Decl.h

@@ -3231,6 +3231,10 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// declared with @TFParameter).
   void getAllTFParameters(SmallVectorImpl<VarDecl *> &result) const;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Get the memberwise initializer of the nominal type, if it exists.
+  ConstructorDecl *getMemberwiseInitializer();
+
 private:
   /// Predicate used to filter StoredPropertyRange.
   struct ToStoredProperty {

lib/AST/Decl.cpp

@@ -3025,6 +3025,29 @@ NominalTypeDecl::getAllTFParameters(SmallVectorImpl<VarDecl *> &result) const {
       result.push_back(member);
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+ConstructorDecl *NominalTypeDecl::getMemberwiseInitializer() {
+  ConstructorDecl *memberwiseInitDecl = nullptr;
+  auto ctorDecls = lookupDirect(DeclBaseName::createConstructor());
+  for (auto decl : ctorDecls) {
+    auto ctorDecl = dyn_cast<ConstructorDecl>(decl);
+    if (!ctorDecl)
+      continue;
+    // Continue if:
+    // - Constructor is not a memberwise initializer.
+    // - Constructor is implicit and takes no arguments, and nominal has no
+    //   stored properties. This is ad-hoc and accepts empty struct
+    //   constructors generated via `TypeChecker::defineDefaultConstructor`.
+    if (!ctorDecl->isMemberwiseInitializer() &&
+        !(getStoredProperties().empty() && ctorDecl->isImplicit() &&
+          ctorDecl->getParameters()->size() == 0))
+      continue;
+    assert(!memberwiseInitDecl && "Memberwise initializer already found");
+    memberwiseInitDecl = ctorDecl;
+  }
+  return memberwiseInitDecl;
+}
+
 GenericTypeDecl::GenericTypeDecl(DeclKind K, DeclContext *DC,
                                  Identifier name, SourceLoc nameLoc,
                                  MutableArrayRef<TypeLoc> inherited,

lib/Sema/DerivedConformanceAdditiveArithmeticVectorNumeric.cpp

@@ -76,22 +76,6 @@ static ValueDecl *getProtocolRequirement(ProtocolDecl *proto, Identifier name) {
   return lookup[0];
 }
 
-// Get memberwise initializer for a nominal type.
-static ConstructorDecl *getMemberwiseInitializer(NominalTypeDecl *nominal) {
-  ConstructorDecl *memberwiseInitDecl = nullptr;
-  for (auto member : nominal->getMembers()) {
-    // Find memberwise initializer.
-    if (!memberwiseInitDecl) {
-      auto initDecl = dyn_cast<ConstructorDecl>(member);
-      if (!initDecl || !initDecl->isMemberwiseInitializer())
-        continue;
-      assert(!memberwiseInitDecl && "Memberwise initializer already found");
-      memberwiseInitDecl = initDecl;
-    }
-  }
-  return memberwiseInitDecl;
-}
-
 // Return the `Scalar` associated type for a ValueDecl if it conforms to
 // `VectorNumeric` in the given context.
 // If the decl does not conform to `VectorNumeric`, return a null `Type`.
@@ -182,7 +166,7 @@ static void deriveBodyMathOperator(AbstractFunctionDecl *funcDecl,
   auto &C = nominal->getASTContext();
 
   // Create memberwise initializer: `Nominal.init(...)`.
-  auto *memberwiseInitDecl = getMemberwiseInitializer(nominal);
+  auto *memberwiseInitDecl = nominal->getMemberwiseInitializer();
   auto *initDRE =
       new (C) DeclRefExpr(memberwiseInitDecl, DeclNameLoc(), /*Implicit*/ true);
   initDRE->setFunctionRefKind(FunctionRefKind::SingleApply);
@@ -361,7 +345,7 @@ static void deriveBodyAdditiveArithmetic_zero(AbstractFunctionDecl *funcDecl) {
   auto *nominal = funcDecl->getDeclContext()->getSelfNominalTypeDecl();
   auto &C = nominal->getASTContext();
 
-  auto *memberwiseInitDecl = getMemberwiseInitializer(nominal);
+  auto *memberwiseInitDecl = nominal->getMemberwiseInitializer();
   auto *initDRE =
       new (C) DeclRefExpr(memberwiseInitDecl, DeclNameLoc(), /*Implicit*/ true);
   initDRE->setFunctionRefKind(FunctionRefKind::SingleApply);
@@ -416,7 +400,7 @@ static ValueDecl *deriveAdditiveArithmetic_zero(DerivedConformance &derived) {
   // The implicit memberwise constructor must be explicitly created so that it
   // can called when synthesizing the `zero` property getter. Normally, the
   // memberwise constructor is synthesized during SILGen, which is too late.
-  if (!getMemberwiseInitializer(nominal)) {
+  if (!nominal->getMemberwiseInitializer()) {
     auto *initDecl = createImplicitConstructor(
         TC, nominal, ImplicitConstructorKind::Memberwise);
     nominal->addMember(initDecl);