AutoDiff: Correctly compute SIL return type in VJPCloner and JVPCloner

Author: slavapestov

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -702,22 +702,16 @@ class JVPCloner::Implementation final
     extractAllElements(origResult, builder, origResults);
 
     // Get and partially apply the differential.
-    auto jvpGenericEnv = jvp->getGenericEnvironment();
-    auto jvpSubstMap = jvpGenericEnv
-                           ? jvpGenericEnv->getForwardingSubstitutionMap()
-                           : jvp->getForwardingSubstitutionMap();
+    auto jvpSubstMap = jvp->getForwardingSubstitutionMap();
     auto *differentialRef = builder.createFunctionRef(loc, &getDifferential());
     auto *differentialPartialApply = builder.createPartialApply(
         loc, differentialRef, jvpSubstMap, {diffStructVal},
         ParameterConvention::Direct_Guaranteed);
 
-    auto differentialType = jvp->getLoweredFunctionType()
-                                ->getResults()
-                                .back()
-                                .getSILStorageInterfaceType();
-    differentialType = differentialType.substGenericArgs(
-        getModule(), jvpSubstMap, TypeExpansionContext::minimal());
-    differentialType = differentialType.subst(getModule(), jvpSubstMap);
+    auto differentialType = jvp->mapTypeIntoContext(
+        jvp->getConventions().getSILType(
+            jvp->getLoweredFunctionType()->getResults().back(),
+            jvp->getTypeExpansionContext()));
     auto differentialFnType = differentialType.castTo<SILFunctionType>();
     auto differentialSubstType =
         differentialPartialApply->getType().castTo<SILFunctionType>();

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -231,10 +231,7 @@ class VJPCloner::Implementation final
     extractAllElements(origResult, Builder, origResults);
 
     // Get and partially apply the pullback.
-    auto vjpGenericEnv = vjp->getGenericEnvironment();
-    auto vjpSubstMap = vjpGenericEnv
-                           ? vjpGenericEnv->getForwardingSubstitutionMap()
-                           : vjp->getForwardingSubstitutionMap();
+    auto vjpSubstMap = vjp->getForwardingSubstitutionMap();
     auto *pullbackRef = Builder.createFunctionRef(loc, pullback);
 
     // Prepare partial application arguments.
@@ -257,13 +254,10 @@ class VJPCloner::Implementation final
     auto *pullbackPartialApply = Builder.createPartialApply(
         loc, pullbackRef, vjpSubstMap, {partialApplyArg},
         ParameterConvention::Direct_Guaranteed);
-    auto pullbackType = vjp->getLoweredFunctionType()
-                            ->getResults()
-                            .back()
-                            .getSILStorageInterfaceType();
-    pullbackType = pullbackType.substGenericArgs(
-        getModule(), vjpSubstMap, TypeExpansionContext::minimal());
-    pullbackType = pullbackType.subst(getModule(), vjpSubstMap);
+    auto pullbackType = vjp->mapTypeIntoContext(
+        vjp->getConventions().getSILType(
+            vjp->getLoweredFunctionType()->getResults().back(),
+            vjp->getTypeExpansionContext()));
     auto pullbackFnType = pullbackType.castTo<SILFunctionType>();
     auto pullbackSubstType =
         pullbackPartialApply->getType().castTo<SILFunctionType>();