Conform Differentiable assoc. types to ElementaryFunctions if possible.

Conform `Differentiable` synthesized associated types to `ElementaryFunctions`
if possible. Similar to existing logic for `VectorProtocol`.

Author: dan-zheng

lib/Sema/DerivedConformanceAdditiveArithmetic.cpp

@@ -183,7 +183,7 @@ static void deriveBodyMathOperator(AbstractFunctionDecl *funcDecl,
     memberOpExprs.push_back(createMemberOpExpr(member));
     memberNames.push_back(member->getName());
   }
-  // Call memberwise initialier with member operator call expressions.
+  // Call memberwise initializer with member operator call expressions.
   auto *callExpr =
       CallExpr::createImplicit(C, initExpr, memberOpExprs, memberNames);
   ASTNode returnStmt = new (C) ReturnStmt(SourceLoc(), callExpr, true);

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -568,12 +568,14 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
   auto diffableType = TypeLoc::withoutLoc(diffableProto->getDeclaredType());
   auto *addArithProto = C.getProtocol(KnownProtocolKind::AdditiveArithmetic);
   auto addArithType = TypeLoc::withoutLoc(addArithProto->getDeclaredType());
+  auto *mathProto = C.getProtocol(KnownProtocolKind::ElementaryFunctions);
+  auto mathType = TypeLoc::withoutLoc(mathProto->getDeclaredType());
   auto *vectorProto = C.getProtocol(KnownProtocolKind::VectorProtocol);
   auto vectorType = TypeLoc::withoutLoc(vectorProto->getDeclaredType());
   auto *kpIterableProto = C.getProtocol(KnownProtocolKind::KeyPathIterable);
   auto kpIterableType = TypeLoc::withoutLoc(kpIterableProto->getDeclaredType());
 
-  SmallVector<TypeLoc, 3> inherited {diffableType};
+  SmallVector<TypeLoc, 3> inherited{diffableType};
 
   // Cache original members and their associated types for later use.
   SmallVector<VarDecl *, 8> diffProperties;
@@ -589,9 +591,16 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
   bool canDeriveAdditiveArithmetic =
       llvm::all_of(diffProperties, [&](VarDecl *vd) {
         return TC.conformsToProtocol(getAssociatedType(vd, parentDC, id),
-                                     addArithProto, parentDC,
-                                     None);
-        });
+                                     addArithProto, parentDC, None);
+      });
+
+  // Associated struct can derive `ElementaryFunctions` if the associated types
+  // of all stored properties conform to `ElementaryFunctions`.
+  bool canDeriveElementaryFunctions =
+      llvm::all_of(diffProperties, [&](VarDecl *vd) {
+        return TC.conformsToProtocol(getAssociatedType(vd, parentDC, id),
+                                     mathProto, parentDC, None);
+      });
 
   // Associated struct can derive `VectorProtocol` if the associated types of
   // all members conform to `VectorProtocol` and share the same
@@ -625,6 +634,10 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
                               None))
       inherited.push_back(kpIterableType);
   }
+  // If all members conform to `ElementaryFunctions`, make the associated struct
+  // conform to `ElementaryFunctions`.
+  if (canDeriveElementaryFunctions)
+    inherited.push_back(mathType);
   // If all members also conform to `VectorProtocol` with the same `Scalar`
   // type, make the associated struct conform to `VectorProtocol` instead of
   // just `AdditiveArithmetic`.

test/AutoDiff/derived_differentiable_properties.swift

@@ -40,7 +40,7 @@ struct TestNoDerivative : Differentiable {
 // CHECK-AST:         var w: Float
 // CHECK-AST:         @noDerivative internal var technicallyDifferentiable: Float
 // CHECK-AST:         internal init(w: Float, technicallyDifferentiable: Float)
-// CHECK-AST:         internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, VectorProtocol
+// CHECK-AST:         internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, ElementaryFunctions, VectorProtocol
 // CHECK-AST:           internal typealias AllDifferentiableVariables = TestNoDerivative.AllDifferentiableVariables
 // CHECK-AST:           internal typealias TangentVector = TestNoDerivative.AllDifferentiableVariables
 // CHECK-AST:         internal typealias TangentVector = TestNoDerivative.AllDifferentiableVariables
@@ -54,7 +54,7 @@ struct TestKeyPathIterable : Differentiable, KeyPathIterable {
 // CHECK-AST:         var w: Float
 // CHECK-AST:         @noDerivative internal var technicallyDifferentiable: Float
 // CHECK-AST:         internal init(w: Float, technicallyDifferentiable: Float)
-// CHECK-AST:         internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, KeyPathIterable, VectorProtocol
+// CHECK-AST:         internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, KeyPathIterable, ElementaryFunctions, VectorProtocol
 // CHECK-AST:           internal typealias AllDifferentiableVariables = TestKeyPathIterable.AllDifferentiableVariables
 // CHECK-AST:           internal typealias TangentVector = TestKeyPathIterable.AllDifferentiableVariables
 // CHECK-AST:         internal typealias TangentVector = TestKeyPathIterable.AllDifferentiableVariables

test/Sema/struct_differentiable.swift

@@ -8,12 +8,18 @@ func assertAllDifferentiableVariablesEqualsTangentVector<T>(_: T.Type)
 // Verify that a type `T` conforms to `AdditiveArithmetic`.
 func assertConformsToAdditiveArithmetic<T>(_: T.Type) where T : AdditiveArithmetic {}
 
+// Verify that a type `T` conforms to `ElementaryFunctions`.
+func assertConformsToElementaryFunctions<T>(_: T.Type) where T : ElementaryFunctions {}
+
 // Verify that a type `T` conforms to `VectorProtocol`.
 func assertConformsToVectorProtocol<T>(_: T.Type) where T : VectorProtocol {}
 
 struct Empty : Differentiable {}
 func testEmpty() {
   assertConformsToAdditiveArithmetic(Empty.AllDifferentiableVariables.self)
+  assertConformsToAdditiveArithmetic(Empty.TangentVector.self)
+  assertConformsToElementaryFunctions(Empty.AllDifferentiableVariables.self)
+  assertConformsToElementaryFunctions(Empty.TangentVector.self)
 }
 
 // Test interaction with `AdditiveArithmetic` derived conformances.
@@ -130,19 +136,35 @@ func testAllMembersAdditiveArithmetic() {
 }
 
 // Test type `AllMembersVectorProtocol` whose members conforms to `VectorProtocol`,
-// in which case we should make `TangentVector` and `TangentVector` conform to
-// `VectorProtocol`.
+// in which case we should make `AllDifferentiableVariables` and `TangentVector`
+// conform to `VectorProtocol`.
 struct MyVector : VectorProtocol, Differentiable {
   var w: Float
   var b: Float
 }
 struct AllMembersVectorProtocol : Differentiable {
-  var w: MyVector
-  var b: MyVector
+  var v1: MyVector
+  var v2: MyVector
 }
 func testAllMembersVectorProtocol() {
+  assertConformsToVectorProtocol(AllMembersVectorProtocol.AllDifferentiableVariables.self)
   assertConformsToVectorProtocol(AllMembersVectorProtocol.TangentVector.self)
-  assertConformsToVectorProtocol(AllMembersVectorProtocol.TangentVector.self)
+}
+
+// Test type `AllMembersElementaryFunctions` whose members conforms to `ElementaryFunctions`,
+// in which case we should make `AllDifferentiableVariables` and `TangentVector`
+// conform to `ElementaryFunctions`.
+struct MyVector2 : ElementaryFunctions, Differentiable {
+  var w: Float
+  var b: Float
+}
+struct AllMembersElementaryFunctions : Differentiable {
+  var v1: MyVector2
+  var v2: MyVector2
+}
+func testAllMembersElementaryFunctions() {
+  assertConformsToElementaryFunctions(AllMembersElementaryFunctions.AllDifferentiableVariables.self)
+  assertConformsToElementaryFunctions(AllMembersElementaryFunctions.TangentVector.self)
 }
 
 // Test type whose properties are not all differentiable.
@@ -154,6 +176,7 @@ struct DifferentiableSubset : Differentiable {
 }
 func testDifferentiableSubset() {
   assertConformsToAdditiveArithmetic(DifferentiableSubset.AllDifferentiableVariables.self)
+  assertConformsToElementaryFunctions(DifferentiableSubset.AllDifferentiableVariables.self)
   assertConformsToVectorProtocol(DifferentiableSubset.AllDifferentiableVariables.self)
   assertAllDifferentiableVariablesEqualsTangentVector(DifferentiableSubset.self)
   _ = DifferentiableSubset.TangentVector(w: 1, b: 1)