[AutoDiff] Support direct init reference differentiation. (swiftlang#30946)

Support `@differentiable` function conversion for `init` references, in
addition to `func` references and literal closures. Minor usability improvement.

Resolves SR-12562.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -1221,7 +1221,7 @@ ERROR(invalid_autoclosure_forwarding,none,
       "add () to forward @autoclosure parameter", ())
 ERROR(invalid_differentiable_function_conversion_expr,none,
       "a '@differentiable%select{|(linear)}0' function can only be formed from "
-      "a reference to a 'func' or a literal closure", (bool))
+      "a reference to a 'func' or 'init' or a literal closure", (bool))
 NOTE(invalid_differentiable_function_conversion_parameter,none,
      "did you mean to take a '%0' closure?", (StringRef))
 ERROR(invalid_autoclosure_pointer_conversion,none,

lib/Sema/CSApply.cpp

@@ -5971,7 +5971,7 @@ maybeDiagnoseUnsupportedDifferentiableConversion(ConstraintSystem &cs,
         semanticExpr = capture->getClosureBody();
       if (isa<ClosureExpr>(semanticExpr)) return;
       if (auto *declRef = dyn_cast<DeclRefExpr>(semanticExpr)) {
-        if (isa<FuncDecl>(declRef->getDecl())) return;
+        if (isa<AbstractFunctionDecl>(declRef->getDecl())) return;
         // If the referenced decl is a function parameter, the user may want
         // to change the declaration to be a '@differentiable' closure. Emit a
         // note with a fix-it.
@@ -6001,13 +6001,16 @@ maybeDiagnoseUnsupportedDifferentiableConversion(ConstraintSystem &cs,
         if (isa<FuncDecl>(memberRef->getMember().getDecl())) return;
       } else if (auto *dotSyntaxCall =
                      dyn_cast<DotSyntaxCallExpr>(semanticExpr)) {
-        Expr *fnExpr = dotSyntaxCall->getFn()->getSemanticsProvidingExpr();
-        while (auto *autoclosureExpr = dyn_cast<AutoClosureExpr>(fnExpr))
-          if (auto *unwrappedFnExpr = autoclosureExpr->getUnwrappedCurryThunkExpr())
-            fnExpr = unwrappedFnExpr;
         // Recurse on the function expression.
+        auto *fnExpr = dotSyntaxCall->getFn()->getSemanticsProvidingExpr();
         maybeDiagnoseFunctionRef(fnExpr);
         return;
+      } else if (auto *autoclosureExpr = dyn_cast<AutoClosureExpr>(semanticExpr)) {
+        // Peer through curry thunks.
+        if (auto *unwrappedFnExpr = autoclosureExpr->getUnwrappedCurryThunkExpr()) {
+          maybeDiagnoseFunctionRef(unwrappedFnExpr);
+          return;
+        }
       }
       ctx.Diags.diagnose(expr->getLoc(),
                          diag::invalid_differentiable_function_conversion_expr,

test/AutoDiff/Sema/differentiable_func_type.swift

@@ -60,19 +60,19 @@ func fakeGradient<T, U: FloatingPoint>(of f: @differentiable (T) -> U) {}
 
 func takesOpaqueClosure(f: @escaping (Float) -> Float) {
   // expected-note @-1 {{did you mean to take a '@differentiable' closure?}} {{38-38=@differentiable }}
-  // expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or a literal closure}}
+  // expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or 'init' or a literal closure}}
   fakeGradient(of: f)
 }
 
 let globalAddOne: (Float) -> Float = { $0 + 1 }
-// expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or a literal closure}}
+// expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or 'init' or a literal closure}}
 fakeGradient(of: globalAddOne)
 
 func someScope() {
   let localAddOne: (Float) -> Float = { $0 + 1 }
-  // expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or a literal closure}}
+  // expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or 'init' or a literal closure}}
   fakeGradient(of: globalAddOne)
-  // expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or a literal closure}}
+  // expected-error @+1 {{a '@differentiable' function can only be formed from a reference to a 'func' or 'init' or a literal closure}}
   fakeGradient(of: localAddOne)
   // The following case is okay during type checking, but will fail in the AD transform.
   fakeGradient { localAddOne($0) }
@@ -95,10 +95,15 @@ func linearToDifferentiable(_ f: @escaping @differentiable(linear) (Float) -> Fl
 }
 
 func differentiableToLinear(_ f: @escaping @differentiable (Float) -> Float) {
-  // expected-error @+1 {{a '@differentiable(linear)' function can only be formed from a reference to a 'func' or a literal closure}}
+  // expected-error @+1 {{a '@differentiable(linear)' function can only be formed from a reference to a 'func' or 'init' or a literal closure}}
   _ = f as @differentiable(linear) (Float) -> Float
 }
 
+struct Struct: Differentiable {
+  var x: Float
+}
+let _: @differentiable (Float) -> Struct = Struct.init
+
 //===----------------------------------------------------------------------===//
 // Parameter selection (@noDerivative)
 //===----------------------------------------------------------------------===//

test/AutoDiff/validation-test/simple_math.swift

@@ -202,6 +202,9 @@ SimpleMathTests.test("StructMemberwiseInitializer") {
     }
   }
 
+  // Test direct `init` reference.
+  expectEqual(10, pullback(at: 4, in: Foo.init)(.init(stored: 10)))
+
   let 𝛁foo = pullback(at: Float(4), in: { input -> Foo in
     let foo = Foo(stored: input)
     let foo2 = foo + foo