Merge pull request #41792 from DougGregor/collapse-thunk-types

Author: DougGregor

include/swift/SIL/TypeLowering.h

@@ -1204,6 +1204,39 @@ CanSILFunctionType getNativeSILFunctionType(
     Optional<SubstitutionMap> reqtSubs = None,
     ProtocolConformanceRef witnessMethodConformance = ProtocolConformanceRef());
 
+/// The thunk kinds used in the differentiation transform.
+enum class DifferentiationThunkKind {
+  /// A reabstraction thunk.
+  ///
+  /// Reabstraction thunks transform a function-typed value to another one with
+  /// different parameter/result abstraction patterns. This is identical to the
+  /// thunks generated by SILGen.
+  Reabstraction,
+
+  /// An index subset thunk.
+  ///
+  /// An index subset thunk is used transform JVP/VJPs into a version that is
+  /// "wrt" fewer differentiation parameters.
+  /// - Differentials of thunked JVPs use zero for non-requested differentiation
+  ///    parameters.
+  /// - Pullbacks of thunked VJPs discard results for non-requested
+  ///   differentiation parameters.
+  IndexSubset
+};
+
+/// Build the type of a function transformation thunk.
+CanSILFunctionType buildSILFunctionThunkType(
+    SILFunction *fn,
+    CanSILFunctionType &sourceType,
+    CanSILFunctionType &expectedType,
+    CanType &inputSubstType,
+    CanType &outputSubstType,
+    GenericEnvironment *&genericEnv,
+    SubstitutionMap &interfaceSubs,
+    CanType &dynamicSelfType,
+    bool withoutActuallyEscaping,
+    Optional<DifferentiationThunkKind> differentiationThunkKind = None);
+
 } // namespace swift
 
 namespace llvm {

include/swift/SILOptimizer/Differentiation/Thunk.h

@@ -47,41 +47,6 @@ class ADContext;
 // moved to a shared location.
 //===----------------------------------------------------------------------===//
 
-/// The thunk kinds used in the differentiation transform.
-enum class DifferentiationThunkKind {
-  /// A reabstraction thunk.
-  ///
-  /// Reabstraction thunks transform a function-typed value to another one with
-  /// different parameter/result abstraction patterns. This is identical to the
-  /// thunks generated by SILGen.
-  Reabstraction,
-
-  /// An index subset thunk.
-  ///
-  /// An index subset thunk is used transform JVP/VJPs into a version that is
-  /// "wrt" fewer differentiation parameters.
-  /// - Differentials of thunked JVPs use zero for non-requested differentiation
-  ///    parameters.
-  /// - Pullbacks of thunked VJPs discard results for non-requested
-  ///   differentiation parameters.
-  IndexSubset
-};
-
-CanGenericSignature buildThunkSignature(SILFunction *fn, bool inheritGenericSig,
-                                        OpenedArchetypeType *openedExistential,
-                                        GenericEnvironment *&genericEnv,
-                                        SubstitutionMap &contextSubs,
-                                        SubstitutionMap &interfaceSubs,
-                                        ArchetypeType *&newArchetype);
-
-/// Build the type of a function transformation thunk.
-CanSILFunctionType buildThunkType(SILFunction *fn,
-                                  CanSILFunctionType &sourceType,
-                                  CanSILFunctionType &expectedType,
-                                  GenericEnvironment *&genericEnv,
-                                  SubstitutionMap &interfaceSubs,
-                                  bool withoutActuallyEscaping,
-                                  DifferentiationThunkKind thunkKind);
 
 /// Get or create a reabstraction thunk from `fromType` to `toType`, to be
 /// called in `caller`.

lib/SIL/IR/SILFunctionType.cpp

@@ -2407,6 +2407,296 @@ CanSILFunctionType swift::getNativeSILFunctionType(
       substConstant, reqtSubs, witnessMethodConformance, None);
 }
 
+/// Build a generic signature and environment for a re-abstraction thunk.
+///
+/// Most thunks share the generic environment with their original function.
+/// The one exception is if the thunk type involves an open existential,
+/// in which case we "promote" the opened existential to a new generic parameter.
+///
+/// \param SGF - the parent function
+/// \param openedExistential - the opened existential to promote to a generic
+//  parameter, if any
+/// \param inheritGenericSig - whether to inherit the generic signature from the
+/// parent function.
+/// \param genericEnv - the new generic environment
+/// \param contextSubs - map old archetypes to new archetypes
+/// \param interfaceSubs - map interface types to old archetypes
+static CanGenericSignature
+buildThunkSignature(SILFunction *fn,
+                    bool inheritGenericSig,
+                    OpenedArchetypeType *openedExistential,
+                    GenericEnvironment *&genericEnv,
+                    SubstitutionMap &contextSubs,
+                    SubstitutionMap &interfaceSubs,
+                    ArchetypeType *&newArchetype) {
+  auto *mod = fn->getModule().getSwiftModule();
+  auto &ctx = mod->getASTContext();
+
+  // If there's no opened existential, we just inherit the generic environment
+  // from the parent function.
+  if (openedExistential == nullptr) {
+    auto genericSig =
+      fn->getLoweredFunctionType()->getInvocationGenericSignature();
+    genericEnv = fn->getGenericEnvironment();
+    interfaceSubs = fn->getForwardingSubstitutionMap();
+    contextSubs = interfaceSubs;
+    return genericSig;
+  }
+
+  // Add the existing generic signature.
+  int depth = 0;
+  GenericSignature baseGenericSig;
+  if (inheritGenericSig) {
+    if (auto genericSig =
+          fn->getLoweredFunctionType()->getInvocationGenericSignature()) {
+      baseGenericSig = genericSig;
+      depth = genericSig.getGenericParams().back()->getDepth() + 1;
+    }
+  }
+
+  // Add a new generic parameter to replace the opened existential.
+  auto *newGenericParam =
+      GenericTypeParamType::get(/*type sequence*/ false, depth, 0, ctx);
+
+  assert(openedExistential->isRoot());
+  auto constraint = openedExistential->getExistentialType();
+  if (auto existential = constraint->getAs<ExistentialType>())
+    constraint = existential->getConstraintType();
+
+  Requirement newRequirement(RequirementKind::Conformance, newGenericParam,
+                             constraint);
+
+  auto genericSig = buildGenericSignature(ctx, baseGenericSig,
+                                          { newGenericParam },
+                                          { newRequirement });
+  genericEnv = genericSig.getGenericEnvironment();
+
+  newArchetype = genericEnv->mapTypeIntoContext(newGenericParam)
+    ->castTo<ArchetypeType>();
+
+  // Calculate substitutions to map the caller's archetypes to the thunk's
+  // archetypes.
+  if (auto calleeGenericSig = fn->getLoweredFunctionType()
+          ->getInvocationGenericSignature()) {
+    contextSubs = SubstitutionMap::get(
+      calleeGenericSig,
+      [&](SubstitutableType *type) -> Type {
+        return genericEnv->mapTypeIntoContext(type);
+      },
+      MakeAbstractConformanceForGenericType());
+  }
+
+  // Calculate substitutions to map interface types to the caller's archetypes.
+  interfaceSubs = SubstitutionMap::get(
+    genericSig,
+    [&](SubstitutableType *type) -> Type {
+      if (type->isEqual(newGenericParam))
+        return openedExistential;
+      return fn->mapTypeIntoContext(type);
+    },
+    MakeAbstractConformanceForGenericType());
+
+  return genericSig.getCanonicalSignature();
+}
+
+/// Build the type of a function transformation thunk.
+CanSILFunctionType swift::buildSILFunctionThunkType(
+    SILFunction *fn,
+    CanSILFunctionType &sourceType,
+    CanSILFunctionType &expectedType,
+    CanType &inputSubstType,
+    CanType &outputSubstType,
+    GenericEnvironment *&genericEnv,
+    SubstitutionMap &interfaceSubs,
+    CanType &dynamicSelfType,
+    bool withoutActuallyEscaping,
+    Optional<DifferentiationThunkKind> differentiationThunkKind) {
+  // We shouldn't be thunking generic types here, and substituted function types
+  // ought to have their substitutions applied before we get here.
+  assert(!expectedType->isPolymorphic() &&
+         !expectedType->getCombinedSubstitutions());
+  assert(!sourceType->isPolymorphic() &&
+         !sourceType->getCombinedSubstitutions());
+
+  // This may inherit @noescape from the expectedType. The @noescape attribute
+  // is only stripped when using this type to materialize a new decl.
+  auto extInfoBuilder = expectedType->getExtInfo().intoBuilder();
+  if (!differentiationThunkKind ||
+      *differentiationThunkKind == DifferentiationThunkKind::Reabstraction ||
+      extInfoBuilder.hasContext()) {
+    // Can't build a reabstraction thunk without context, so we require
+    // ownership semantics on the result type.
+    assert(expectedType->getExtInfo().hasContext());
+
+    extInfoBuilder = extInfoBuilder.withRepresentation(
+          SILFunctionType::Representation::Thin);
+  }
+
+  if (withoutActuallyEscaping)
+    extInfoBuilder = extInfoBuilder.withNoEscape(false);
+
+  // Does the thunk type involve archetypes other than opened existentials?
+  bool hasArchetypes = false;
+  // Does the thunk type involve an open existential type?
+  CanOpenedArchetypeType openedExistential;
+  auto archetypeVisitor = [&](CanType t) {
+    if (auto archetypeTy = dyn_cast<ArchetypeType>(t)) {
+      if (auto opened = dyn_cast<OpenedArchetypeType>(archetypeTy)) {
+        const auto root = cast<OpenedArchetypeType>(CanType(opened->getRoot()));
+        assert((openedExistential == CanArchetypeType() ||
+                openedExistential == root) &&
+               "one too many open existentials");
+        openedExistential = root;
+      } else {
+        hasArchetypes = true;
+      }
+    }
+  };
+
+  // Use the generic signature from the context if the thunk involves
+  // generic parameters.
+  CanGenericSignature genericSig;
+  SubstitutionMap contextSubs;
+  ArchetypeType *newArchetype = nullptr;
+
+  if (expectedType->hasArchetype() || sourceType->hasArchetype()) {
+    expectedType.visit(archetypeVisitor);
+    sourceType.visit(archetypeVisitor);
+
+    genericSig = buildThunkSignature(fn,
+                                     hasArchetypes,
+                                     openedExistential,
+                                     genericEnv,
+                                     contextSubs,
+                                     interfaceSubs,
+                                     newArchetype);
+  }
+
+  auto substTypeHelper = [&](SubstitutableType *type) -> Type {
+    if (CanType(type) == openedExistential)
+      return newArchetype;
+
+    // If a nested archetype is rooted on our opened existential, fail:
+    // Type::subst attempts to substitute the parent of a nested archetype
+    // only if it fails to find a replacement for the nested one.
+    if (auto *opened = dyn_cast<OpenedArchetypeType>(type)) {
+      if (openedExistential->isEqual(opened->getRoot())) {
+        return nullptr;
+      }
+    }
+
+    return Type(type).subst(contextSubs);
+  };
+  auto substConformanceHelper =
+    LookUpConformanceInSubstitutionMap(contextSubs);
+
+  // Utility function to apply contextSubs, and also replace the
+  // opened existential with the new archetype.
+  auto substFormalTypeIntoThunkContext =
+      [&](CanType t) -> CanType {
+    return t.subst(substTypeHelper, substConformanceHelper)
+               ->getCanonicalType();
+  };
+  auto substLoweredTypeIntoThunkContext =
+      [&](CanSILFunctionType t) -> CanSILFunctionType {
+    return SILType::getPrimitiveObjectType(t)
+             .subst(fn->getModule(), substTypeHelper, substConformanceHelper)
+             .castTo<SILFunctionType>();
+  };
+
+  sourceType = substLoweredTypeIntoThunkContext(sourceType);
+  expectedType = substLoweredTypeIntoThunkContext(expectedType);
+
+  bool hasDynamicSelf = false;
+
+  if (inputSubstType) {
+    inputSubstType = substFormalTypeIntoThunkContext(inputSubstType);
+    hasDynamicSelf |= inputSubstType->hasDynamicSelfType();
+  }
+
+  if (outputSubstType) {
+    outputSubstType = substFormalTypeIntoThunkContext(outputSubstType);
+    hasDynamicSelf |= outputSubstType->hasDynamicSelfType();
+  }
+
+  hasDynamicSelf |= sourceType->hasDynamicSelfType();
+  hasDynamicSelf |= expectedType->hasDynamicSelfType();
+
+  // If our parent function was pseudogeneric, this thunk must also be
+  // pseudogeneric, since we have no way to pass generic parameters.
+  if (genericSig)
+    if (fn->getLoweredFunctionType()->isPseudogeneric())
+      extInfoBuilder = extInfoBuilder.withIsPseudogeneric();
+
+  // Add the function type as the parameter.
+  auto contextConvention =
+      fn->getTypeLowering(sourceType).isTrivial()
+          ? ParameterConvention::Direct_Unowned
+          : ParameterConvention::Direct_Guaranteed;
+  SmallVector<SILParameterInfo, 4> params;
+  params.append(expectedType->getParameters().begin(),
+                expectedType->getParameters().end());
+
+  if (!differentiationThunkKind ||
+      *differentiationThunkKind == DifferentiationThunkKind::Reabstraction) {
+    params.push_back({sourceType,
+                      sourceType->getExtInfo().hasContext()
+                        ? contextConvention
+                        : ParameterConvention::Direct_Unowned});
+  }
+
+  // If this thunk involves DynamicSelfType in any way, add a capture for it
+  // in case we need to recover metadata.
+  if (hasDynamicSelf) {
+    dynamicSelfType = fn->getDynamicSelfMetadata()->getType().getASTType();
+    if (!isa<MetatypeType>(dynamicSelfType)) {
+      dynamicSelfType = CanMetatypeType::get(dynamicSelfType,
+                                             MetatypeRepresentation::Thick);
+    }
+    params.push_back({dynamicSelfType, ParameterConvention::Direct_Unowned});
+  }
+
+  auto mapTypeOutOfContext = [&](CanType type) -> CanType {
+    return type->mapTypeOutOfContext()->getCanonicalType(genericSig);
+  };
+
+  // Map the parameter and expected types out of context to get the interface
+  // type of the thunk.
+  SmallVector<SILParameterInfo, 4> interfaceParams;
+  interfaceParams.reserve(params.size());
+  for (auto &param : params) {
+    auto interfaceParam = param.map(mapTypeOutOfContext);
+    interfaceParams.push_back(interfaceParam);
+  }
+
+  SmallVector<SILYieldInfo, 4> interfaceYields;
+  for (auto &yield : expectedType->getYields()) {
+    auto interfaceYield = yield.map(mapTypeOutOfContext);
+    interfaceYields.push_back(interfaceYield);
+  }
+
+  SmallVector<SILResultInfo, 4> interfaceResults;
+  for (auto &result : expectedType->getResults()) {
+    auto interfaceResult = result.map(mapTypeOutOfContext);
+    interfaceResults.push_back(interfaceResult);
+  }
+
+  Optional<SILResultInfo> interfaceErrorResult;
+  if (expectedType->hasErrorResult()) {
+    auto errorResult = expectedType->getErrorResult();
+    interfaceErrorResult = errorResult.map(mapTypeOutOfContext);;
+  }
+
+  // The type of the thunk function.
+  return SILFunctionType::get(
+      genericSig, extInfoBuilder.build(), expectedType->getCoroutineKind(),
+      ParameterConvention::Direct_Unowned, interfaceParams, interfaceYields,
+      interfaceResults, interfaceErrorResult,
+      expectedType->getPatternSubstitutions(), SubstitutionMap(),
+      fn->getASTContext());
+
+}
+
 //===----------------------------------------------------------------------===//
 //                          Foreign SILFunctionTypes
 //===----------------------------------------------------------------------===//