[AutoDiff] Fix incorrectly sorted members of synthesized 'TangentVector'. (#36154)

Do not mark synthesized `TangentVector` members as synthesized so that the type checker won't sort them. We would like tangent vector members to have the same order as the properties in the parent declaration.

Also add `typealias TangentVector = Self` to the synthesized `TangentVector` so that it will not need its own `Differentiable` conformance derivation.

Resolves SR-14241 / rdar://74659803.

Author: rxwei

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -40,7 +40,7 @@ using namespace swift;
 /// If the given property is a `var`, return true because `move(by:)` can be
 /// invoked regardless.  Otherwise, return true if and only if the property's
 /// type's 'Differentiable.move(by:)' witness is non-mutating.
-static bool canInvokeMoveAlongOnProperty(
+static bool canInvokeMoveByOnProperty(
     VarDecl *vd, ProtocolConformanceRef diffableConformance) {
   assert(diffableConformance && "Property must conform to 'Differentiable'");
   // `var` always supports `move(by:)` since it is mutable.
@@ -64,7 +64,7 @@ static void
 getStoredPropertiesForDifferentiation(
     NominalTypeDecl *nominal, DeclContext *DC,
     SmallVectorImpl<VarDecl *> &result,
-    bool includeLetPropertiesWithNonmutatingMoveAlong = false) {
+    bool includeLetPropertiesWithNonmutatingMoveBy = false) {
   auto &C = nominal->getASTContext();
   auto *diffableProto = C.getProtocol(KnownProtocolKind::Differentiable);
   for (auto *vd : nominal->getStoredProperties()) {
@@ -90,8 +90,8 @@ getStoredPropertiesForDifferentiation(
     // Skip `let` stored properties with a mutating `move(by:)` if requested.
     // `mutating func move(by:)` cannot be synthesized to update `let`
     // properties.
-    if (!includeLetPropertiesWithNonmutatingMoveAlong && 
-        !canInvokeMoveAlongOnProperty(vd, conformance))
+    if (!includeLetPropertiesWithNonmutatingMoveBy && 
+        !canInvokeMoveByOnProperty(vd, conformance))
       continue;
     result.push_back(vd);
   }
@@ -450,9 +450,9 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
     auto *tangentProperty = new (C) VarDecl(
         member->isStatic(), member->getIntroducer(),
         /*NameLoc*/ SourceLoc(), member->getName(), structDecl);
-    tangentProperty->setSynthesized();
-    // Note: `tangentProperty` is not marked as implicit here, because that
-    // incorrectly affects memberwise initializer synthesis.
+    // Note: `tangentProperty` is not marked as implicit or synthesized here,
+    // because that incorrectly affects memberwise initializer synthesis and
+    // causes the type checker to not guarantee the order of these members.
     auto memberContextualType =
         parentDC->mapTypeIntoContext(member->getValueInterfaceType());
     auto memberTanType =
@@ -507,20 +507,27 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
     }
   }
 
-  // If nominal type is `@_fixed_layout`, also mark `TangentVector` struct as
-  // `@_fixed_layout`.
-  if (nominal->getAttrs().hasAttribute<FixedLayoutAttr>())
-    addFixedLayoutAttr(structDecl);
-
-  // If nominal type is `@frozen`, also mark `TangentVector` struct as
-  // `@frozen`.
+  // If nominal type is `@frozen`, also mark `TangentVector` struct.
   if (nominal->getAttrs().hasAttribute<FrozenAttr>())
     structDecl->getAttrs().add(new (C) FrozenAttr(/*implicit*/ true));
-
-  // If nominal type is `@usableFromInline`, also mark `TangentVector` struct as
-  // `@usableFromInline`.
-  if (nominal->getAttrs().hasAttribute<UsableFromInlineAttr>())
+  
+  // Add `typealias TangentVector = Self` so that the `TangentVector` itself
+  // won't need its own conformance derivation.
+  auto *tangentEqualsSelfAlias = new (C) TypeAliasDecl(
+      SourceLoc(), SourceLoc(), C.Id_TangentVector, SourceLoc(),
+      /*GenericParams*/ nullptr, structDecl);
+  tangentEqualsSelfAlias->setUnderlyingType(structDecl->getSelfTypeInContext());
+  tangentEqualsSelfAlias->setAccess(structDecl->getFormalAccess());
+  tangentEqualsSelfAlias->setImplicit();
+  tangentEqualsSelfAlias->setSynthesized();
+  structDecl->addMember(tangentEqualsSelfAlias);
+
+  // If nominal type is `@usableFromInline`, also mark `TangentVector` struct.
+  if (nominal->getAttrs().hasAttribute<UsableFromInlineAttr>()) {
     structDecl->getAttrs().add(new (C) UsableFromInlineAttr(/*implicit*/ true));
+    tangentEqualsSelfAlias->getAttrs().add(
+        new (C) UsableFromInlineAttr(/*implicit*/ true));
+  }
 
   // The implicit memberwise constructor must be explicitly created so that it
   // can called in `AdditiveArithmetic` and `Differentiable` methods. Normally,
@@ -593,7 +600,7 @@ static void checkAndDiagnoseImplicitNoDerivative(ASTContext &Context,
         TypeChecker::conformsToProtocol(varType, diffableProto, nominal);
     // If stored property should not be diagnosed, continue.
     if (diffableConformance && 
-        canInvokeMoveAlongOnProperty(vd, diffableConformance))
+        canInvokeMoveByOnProperty(vd, diffableConformance))
       continue;
     // Otherwise, add an implicit `@noDerivative` attribute.
     vd->getAttrs().add(new (Context) NoDerivativeAttr(/*Implicit*/ true));

test/AutoDiff/Sema/DerivedConformances/derived_differentiable.swift

@@ -1,4 +1,5 @@
 // RUN: %target-swift-frontend -print-ast %s | %FileCheck %s --check-prefix=CHECK-AST
+// RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s --check-prefix=CHECK-SIL
 
 import _Differentiation
 
@@ -156,3 +157,24 @@ extension TangentVectorP where Self == StructWithTangentVectorConstrained.Tangen
 
 // CHECK-AST-LABEL: internal struct StructWithTangentVectorConstrained : TangentVectorConstrained {
 // CHECK-AST:   internal struct TangentVector : {{(TangentVectorP, Differentiable, AdditiveArithmetic)|(TangentVectorP, AdditiveArithmetic, Differentiable)|(Differentiable, TangentVectorP, AdditiveArithmetic)|(AdditiveArithmetic, TangentVectorP, Differentiable)|(Differentiable, AdditiveArithmetic, TangentVectorP)|(AdditiveArithmetic, Differentiable, TangentVectorP)}} {
+
+public struct SR14241Struct: Differentiable {
+  public var simd: [Float]
+  public var scalar: Float
+}
+
+// CHECK-AST-LABEL: public struct SR14241Struct : Differentiable {
+// CHECK-AST: public var simd: [Float]
+// CHECK-AST: public var scalar: Float
+// CHECK-AST: struct TangentVector : AdditiveArithmetic, Differentiable {
+// CHECK-AST:   var simd: Array<Float>.TangentVector
+// CHECK-AST:   var scalar: Float
+
+// CHECK-SIL-LABEL: public struct SR14241Struct : Differentiable {
+// CHECK-SIL: @differentiable(reverse, wrt: self)
+// CHECK-SIL: @_hasStorage public var simd: [Float] { get set }
+// CHECK-SIL: @differentiable(reverse, wrt: self)
+// CHECK-SIL: @_hasStorage public var scalar: Float { get set }
+// CHECK-SIL: struct TangentVector : AdditiveArithmetic, Differentiable {
+// CHECK-SIL:   @_hasStorage var simd: Array<Float>.DifferentiableView { get set }
+// CHECK-SIL:   @_hasStorage var scalar: Float { get set }