[AutoDiff] Derive Differentiable.zeroTangentVectorInitializer.

docs/DifferentiableProgramming.md

@@ -1079,11 +1079,6 @@ public extension Differentiable where Self == TangentVector {
     mutating func move(along direction: TangentVector) {
         self += direction
     }
-
-    @noDerivative
-    var zeroTangentVectorInitializer: () -> TangentVector {
-        { .zero }
-    }
 }
 ```
 
@@ -1144,8 +1139,8 @@ extension Array: Differentiable where Element: Differentiable {
 
     @noDerivative
     public var zeroTangentVectorInitializer: () -> TangentVector {
-        { [count = self.count] in
-            TangentVector(Array(repeating: .zero, count: count))
+        { [zeroInits = map(\.zeroTangentVectorInitializer)] in
+            TangentVector(zeroInits.map { $0() })
         }
     }
 }
@@ -1238,8 +1233,15 @@ the same effective access level as their corresponding original properties.
 
 A `move(along:)` method is synthesized with a body that calls `move(along:)` for
 each pair of the original property and its corresponding property in
-`TangentVector`. Similarly, `zeroTangentVector` is synthesized to return a
-tangent vector that consists of each stored property's `zeroTangentVector`.
+`TangentVector`.
+
+Similarly, when memberwise derivation is possible,
+`zeroTangentVectorInitializer` is synthesized to return a closure that captures
+and calls each stored property's `zeroTangentVectorInitializer` closure.
+When memberwise derivation is not possible (e.g. for custom user-defined
+`TangentVector` types), `zeroTangentVectorInitializer` is synthesized as a
+`{ TangentVector.zero }` closure.
+
 Here's an example:
 
 ```swift
@@ -1251,14 +1253,17 @@ struct Foo<T: Differentiable, U: Differentiable>: @memberwise Differentiable {
     @noDerivative let helperVariable: T
 
     // The compiler synthesizes:
+    //
     //     struct TangentVector: Differentiable, AdditiveArithmetic {
     //         var x: T.TangentVector
     //         var y: U.TangentVector
     //     }
+    //
     //     mutating func move(along direction: TangentVector) {
     //         x.move(along: direction.x)
     //         y.move(along: direction.y)
     //     }
+    //
     //     @noDerivative
     //     var zeroTangentVectorInitializer: () -> TangentVector {
     //         { [xTanInit = x.zeroTangentVectorInitializer,
@@ -1278,16 +1283,25 @@ properties are declared to conform to `AdditiveArithmetic`. There are no
 `@noDerivative` stored properties.
 
 In these cases, the compiler will make `TangentVector` be a type alias for Self.
-Method `move(along:)` and property `zeroTangentVector` will not be synthesized
-because a default implementation already exists.
+Method `move(along:)` will not be synthesized because a default implementation
+already exists.
 
 ```swift
 struct Point<T: Real>: @memberwise Differentiable, @memberwise AdditiveArithmetic {
     // `x` and `y` are the "differentiation properties".
     var x, y: T
 
     // The compiler synthesizes:
+    //
     //     typealias TangentVector = Self
+    //
+    //     @noDerivative
+    //     var zeroTangentVectorInitializer: () -> TangentVector {
+    //         { [xTanInit = x.zeroTangentVectorInitializer,
+    //            yTanInit = y.zeroTangentVectorInitializer] in
+    //             TangentVector(x: xTanInit(), y: yTanInit())
+    //         }
+    //     }
 }
 ```
 

include/swift/AST/KnownIdentifiers.def

@@ -223,6 +223,7 @@ IDENTIFIER(move)
 IDENTIFIER(pullback)
 IDENTIFIER(TangentVector)
 IDENTIFIER(zero)
+IDENTIFIER(zeroTangentVectorInitializer)
 
 #undef IDENTIFIER
 #undef IDENTIFIER_

lib/Sema/CodeSynthesis.cpp

@@ -1414,3 +1414,22 @@ void swift::addFixedLayoutAttr(NominalTypeDecl *nominal) {
   // Add `@_fixed_layout` to the nominal.
   nominal->getAttrs().add(new (C) FixedLayoutAttr(/*Implicit*/ true));
 }
+
+Expr *DiscriminatorFinder::walkToExprPost(Expr *E) {
+  auto *ACE = dyn_cast<AbstractClosureExpr>(E);
+  if (!ACE)
+    return E;
+
+  unsigned Discriminator = ACE->getDiscriminator();
+  assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
+         "Existing closures should have valid discriminators");
+  if (Discriminator >= NextDiscriminator)
+    NextDiscriminator = Discriminator + 1;
+  return E;
+}
+
+unsigned DiscriminatorFinder::getNextDiscriminator() {
+  if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
+    llvm::report_fatal_error("Out of valid closure discriminators");
+  return NextDiscriminator++;
+}

lib/Sema/CodeSynthesis.h

@@ -18,6 +18,7 @@
 #ifndef SWIFT_TYPECHECKING_CODESYNTHESIS_H
 #define SWIFT_TYPECHECKING_CODESYNTHESIS_H
 
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/ForeignErrorConvention.h"
 #include "swift/Basic/ExternalUnion.h"
 #include "swift/Basic/LLVM.h"
@@ -75,6 +76,20 @@ bool hasLetStoredPropertyWithInitialValue(NominalTypeDecl *nominal);
 /// Add `@_fixed_layout` attribute to the nominal type, if possible.
 void addFixedLayoutAttr(NominalTypeDecl *nominal);
 
+/// Find available closure discriminators.
+///
+/// The parser typically takes care of assigning unique discriminators to
+/// closures, but the parser is unavailable during semantic analysis.
+class DiscriminatorFinder : public ASTWalker {
+  unsigned NextDiscriminator = 0;
+
+public:
+  Expr *walkToExprPost(Expr *E) override;
+
+  // Get the next available closure discriminator.
+  unsigned getNextDiscriminator();
+};
+
 } // end namespace swift
 
 #endif

lib/Sema/DebuggerTestingTransform.cpp

@@ -15,6 +15,7 @@
 ///
 //===----------------------------------------------------------------------===//
 
+#include "CodeSynthesis.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTNode.h"
 #include "swift/AST/ASTWalker.h"
@@ -33,35 +34,6 @@ using namespace swift;
 
 namespace {
 
-/// Find available closure discriminators.
-///
-/// The parser typically takes care of assigning unique discriminators to
-/// closures, but the parser is unavailable to this transform.
-class DiscriminatorFinder : public ASTWalker {
-  unsigned NextDiscriminator = 0;
-
-public:
-  Expr *walkToExprPost(Expr *E) override {
-    auto *ACE = dyn_cast<AbstractClosureExpr>(E);
-    if (!ACE)
-      return E;
-
-    unsigned Discriminator = ACE->getDiscriminator();
-    assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
-           "Existing closures should have valid discriminators");
-    if (Discriminator >= NextDiscriminator)
-      NextDiscriminator = Discriminator + 1;
-    return E;
-  }
-
-  // Get the next available closure discriminator.
-  unsigned getNextDiscriminator() {
-    if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
-      llvm::report_fatal_error("Out of valid closure discriminators");
-    return NextDiscriminator++;
-  }
-};
-
 /// Instrument decls with sanity-checks which the debugger can evaluate.
 class DebuggerTestingTransform : public ASTWalker {
   ASTContext &Ctx;