[AutoDiff] Add missing withoutDerivative(at:) fix-its. (swiftlang#33660)

Add `withoutDerivative(at:)` fix-its for errors regarding non-differentiable
arguments and results.

Author: dan-zheng

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -553,9 +553,14 @@ class JVPCloner::Implementation final
               if (!originalFnTy->getParameters()[paramIndex]
                        .getSILStorageInterfaceType()
                        .isDifferentiable(getModule())) {
-                context.emitNondifferentiabilityError(
-                    ai->getArgumentsWithoutIndirectResults()[paramIndex],
-                    invoker, diag::autodiff_nondifferentiable_argument);
+                auto arg = ai->getArgumentsWithoutIndirectResults()[paramIndex];
+                auto startLoc = arg.getLoc().getStartSourceLoc();
+                auto endLoc = arg.getLoc().getEndSourceLoc();
+                context
+                    .emitNondifferentiabilityError(
+                        arg, invoker, diag::autodiff_nondifferentiable_argument)
+                    .fixItInsert(startLoc, "withoutDerivative(at: ")
+                    .fixItInsertAfter(endLoc, ")");
                 errorOccurred = true;
                 return true;
               }
@@ -573,9 +578,14 @@ class JVPCloner::Implementation final
                                          .getSILStorageInterfaceType();
               }
               if (!remappedResultType.isDifferentiable(getModule())) {
-                context.emitNondifferentiabilityError(
-                    origCallee, invoker,
-                    diag::autodiff_nondifferentiable_result);
+                auto startLoc = ai->getLoc().getStartSourceLoc();
+                auto endLoc = ai->getLoc().getEndSourceLoc();
+                context
+                    .emitNondifferentiabilityError(
+                        origCallee, invoker,
+                        diag::autodiff_nondifferentiable_result)
+                    .fixItInsert(startLoc, "withoutDerivative(at: ")
+                    .fixItInsertAfter(endLoc, ")");
                 errorOccurred = true;
                 return true;
               }

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -457,9 +457,14 @@ class VJPCloner::Implementation final
             if (!originalFnTy->getParameters()[paramIndex]
                      .getSILStorageInterfaceType()
                      .isDifferentiable(getModule())) {
-              context.emitNondifferentiabilityError(
-                  ai->getArgumentsWithoutIndirectResults()[paramIndex], invoker,
-                  diag::autodiff_nondifferentiable_argument);
+              auto arg = ai->getArgumentsWithoutIndirectResults()[paramIndex];
+              auto startLoc = arg.getLoc().getStartSourceLoc();
+              auto endLoc = arg.getLoc().getEndSourceLoc();
+              context
+                  .emitNondifferentiabilityError(
+                      arg, invoker, diag::autodiff_nondifferentiable_argument)
+                  .fixItInsert(startLoc, "withoutDerivative(at: ")
+                  .fixItInsertAfter(endLoc, ")");
               errorOccurred = true;
               return true;
             }
@@ -477,8 +482,14 @@ class VJPCloner::Implementation final
                                        .getSILStorageInterfaceType();
             }
             if (!remappedResultType.isDifferentiable(getModule())) {
-              context.emitNondifferentiabilityError(
-                  origCallee, invoker, diag::autodiff_nondifferentiable_result);
+              auto startLoc = ai->getLoc().getStartSourceLoc();
+              auto endLoc = ai->getLoc().getEndSourceLoc();
+              context
+                  .emitNondifferentiabilityError(
+                      origCallee, invoker,
+                      diag::autodiff_nondifferentiable_result)
+                  .fixItInsert(startLoc, "withoutDerivative(at: ")
+                  .fixItInsertAfter(endLoc, ")");
               errorOccurred = true;
               return true;
             }

test/AutoDiff/SILOptimizer/differentiation_control_flow_diagnostics.swift

@@ -171,7 +171,7 @@ func loop_array(_ array: [Float]) -> Float {
   var result: Float = 1
   // TODO(TF-957): Improve non-differentiability errors for for-in loops
   // (`Collection.makeIterator` and `IteratorProtocol.next`).
-  // expected-note @+1 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}}
+  // expected-note @+1 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}} {{12-12=withoutDerivative(at: }} {{17-17=)}}
   for x in array {
     result = result * x
   }

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -301,14 +301,14 @@ struct TF_687<T> : Differentiable {
   }
 }
 // expected-error @+2 {{function is not differentiable}}
-// expected-note @+1 {{cannot differentiate through a non-differentiable argument; do you want to use 'withoutDerivative(at:)'?}}
+// expected-note @+1 {{cannot differentiate through a non-differentiable argument; do you want to use 'withoutDerivative(at:)'?}} {{78-78=withoutDerivative(at: }} {{79-79=)}}
 let _: @differentiable (Float) -> TF_687<Any> = { x in TF_687<Any>(x, dummy: x) }
 
 // expected-error @+1 {{function is not differentiable}}
 @differentiable
 // expected-note @+1 {{when differentiating this function definition}}
 func roundingGivesError(x: Float) -> Float {
-  // expected-note @+1 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}}
+  // expected-note @+1 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}} {{16-16=withoutDerivative(at: }} {{22-22=)}}
   return Float(Int(x))
 }
 
@@ -688,7 +688,7 @@ func differentiableProjectedValueAccess(_ s: Struct) -> Float {
 // expected-note @+2 {{when differentiating this function definition}}
 @differentiable
 func projectedValueAccess(_ s: Struct) -> Float {
-  // expected-note @+1 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}}
+  // expected-note @+1 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}} {{3-3=withoutDerivative(at: }} {{7-7=)}}
   s.$y.wrappedValue
 }
 
@@ -714,7 +714,7 @@ func modify(_ s: Struct, _ x: Float) -> Float {
 func tupleArrayLiteralInitialization(_ x: Float, _ y: Float) -> Float {
   // `Array<(Float, Float)>` does not conform to `Differentiable`.
   let array = [(x * y, x * y)]
-  // expected-note @+1 {{cannot differentiate through a non-differentiable argument; do you want to use 'withoutDerivative(at:)'?}}
+  // expected-note @+1 {{cannot differentiate through a non-differentiable argument; do you want to use 'withoutDerivative(at:)'?}} {{10-10=withoutDerivative(at: }} {{15-15=)}}
   return array[0].0
 }