From d888f8d2d0dc308d6364c914e25c89b7b950a371 Mon Sep 17 00:00:00 2001
From: Marc Rasi <marcrasi@google.com>
Date: Sun, 5 Apr 2020 16:30:38 -0700
Subject: [PATCH] [AutoDiff] fix SR-12493

We simply needed to upstream `TypeSubstCloner::visitDifferentiableFunctionExtractInst`.

The code has a detailed comment explaining what it does.
---
 include/swift/SIL/TypeSubstCloner.h           | 57 +++++++++++++++++++
 ...function-extract-subst-function-type.swift |  2 +-
 .../stdlib/differential_operators.swift       |  3 -
 3 files changed, 58 insertions(+), 4 deletions(-)
 rename test/AutoDiff/{compiler_crashers => compiler_crashers_fixed}/sr12493-differentiable-function-extract-subst-function-type.swift (91%)

diff --git a/include/swift/SIL/TypeSubstCloner.h b/include/swift/SIL/TypeSubstCloner.h
index 2f4427a34bcbf..b908df726504b 100644
--- a/include/swift/SIL/TypeSubstCloner.h
+++ b/include/swift/SIL/TypeSubstCloner.h
@@ -313,6 +313,63 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     super::visitDestroyValueInst(Destroy);
   }
 
+  void visitDifferentiableFunctionExtractInst(
+      DifferentiableFunctionExtractInst *dfei) {
+    // If the extractee is the original function, do regular cloning.
+    if (dfei->getExtractee() ==
+        NormalDifferentiableFunctionTypeComponent::Original) {
+      super::visitDifferentiableFunctionExtractInst(dfei);
+      return;
+    }
+    // If the extractee is a derivative function, check whether the *remapped
+    // derivative function type* (BC) is equal to the *derivative remapped
+    // function type* (AD).
+    //
+    // +----------------+        remap       +-------------------------+
+    // | orig.  fn type |  -------(A)------> | remapped orig.  fn type |
+    // +----------------+                    +-------------------------+
+    //         |                                                |
+    //    (B, SILGen)   getAutoDiffDerivativeFunctionType   (D, here)
+    //         V                                                V
+    // +----------------+        remap       +-------------------------+
+    // | deriv. fn type |  -------(C)------> | remapped deriv. fn type |
+    // +----------------+                    +-------------------------+
+    //
+    // (AD) does not always commute with (BC):
+    // - (AD) is the result of remapping, then computing the derivative type.
+    //   This is the default cloning behavior, but may break invariants in the
+    //   initial SIL generated by SILGen.
+    // - (BC) is the result of computing the derivative type (SILGen), then
+    //   remapping. This is the expected type, preserving invariants from
+    //   earlier transforms.
+    //
+    // If (AD) is not equal to (BC), use (BC) as the explicit type.
+    SILType remappedOrigType = getOpType(dfei->getOperand()->getType());
+    auto remappedOrigFnType = remappedOrigType.castTo<SILFunctionType>();
+    auto derivativeRemappedFnType =
+        remappedOrigFnType
+            ->getAutoDiffDerivativeFunctionType(
+                remappedOrigFnType->getDifferentiabilityParameterIndices(),
+                /*resultIndex*/ 0, dfei->getDerivativeFunctionKind(),
+                getBuilder().getModule().Types,
+                LookUpConformanceInModule(SwiftMod))
+            ->getWithoutDifferentiability();
+    SILType remappedDerivativeFnType = getOpType(dfei->getType());
+    // If remapped derivative type and derivative remapped type are equal, do
+    // regular cloning.
+    if (SILType::getPrimitiveObjectType(derivativeRemappedFnType) ==
+        remappedDerivativeFnType) {
+      super::visitDifferentiableFunctionExtractInst(dfei);
+      return;
+    }
+    // Otherwise, explicitly use the remapped derivative type.
+    recordClonedInstruction(
+        dfei,
+        getBuilder().createDifferentiableFunctionExtract(
+            getOpLocation(dfei->getLoc()), dfei->getExtractee(),
+            getOpValue(dfei->getOperand()), remappedDerivativeFnType));
+  }
+
   /// One abstract function in the debug info can only have one set of variables
   /// and types. This function determines whether applying the substitutions in
   /// \p SubsMap on the generic signature \p Sig will change the generic type
diff --git a/test/AutoDiff/compiler_crashers/sr12493-differentiable-function-extract-subst-function-type.swift b/test/AutoDiff/compiler_crashers_fixed/sr12493-differentiable-function-extract-subst-function-type.swift
similarity index 91%
rename from test/AutoDiff/compiler_crashers/sr12493-differentiable-function-extract-subst-function-type.swift
rename to test/AutoDiff/compiler_crashers_fixed/sr12493-differentiable-function-extract-subst-function-type.swift
index 6aec648b27081..98952bf453ed4 100644
--- a/test/AutoDiff/compiler_crashers/sr12493-differentiable-function-extract-subst-function-type.swift
+++ b/test/AutoDiff/compiler_crashers_fixed/sr12493-differentiable-function-extract-subst-function-type.swift
@@ -1,4 +1,4 @@
-// RUN: not %target-build-swift -O %s
+// RUN: %target-build-swift -O %s
 
 // SR-12493: SIL verification error regarding substituted function types and
 // `differentiable_function_extract` instruction. Occurs only with `-O`.
diff --git a/test/AutoDiff/stdlib/differential_operators.swift b/test/AutoDiff/stdlib/differential_operators.swift
index 22a7262ac0ada..5376a5d9adb87 100644
--- a/test/AutoDiff/stdlib/differential_operators.swift
+++ b/test/AutoDiff/stdlib/differential_operators.swift
@@ -4,9 +4,6 @@
 // RUN: %target-run %t/differential_operators
 // REQUIRES: executable_test
 
-// FIXME(SR-12493): Disable test for `-O` due to SIL verification error.
-// REQUIRES: swift_test_mode_optimize_none
-
 import _Differentiation
 
 import StdlibUnittest