Merge pull request #73363 from slavapestov/pack-expansion-closures-part-2

Extract common code for building opened existential and element signatures

Author: slavapestov

include/swift/AST/GenericSignature.h

@@ -197,6 +197,11 @@ class GenericSignature {
   /// array of the generic parameters for the innermost generic type.
   ArrayRef<GenericTypeParamType *> getInnermostGenericParams() const;
 
+  /// Returns the depth that a generic parameter at the next level of
+  /// nesting would have. This is zero for the empty signature,
+  /// and one plus the depth of the final generic parameter otherwise.
+  unsigned getNextDepth() const;
+
   /// Retrieve the requirements.
   ArrayRef<Requirement> getRequirements() const;
 
@@ -307,6 +312,9 @@ class alignas(1 << TypeAlignInBits) GenericSignatureImpl final
     return Mem;
   }
 
+  /// Returns the depth of the last generic parameter.
+  unsigned getMaxDepth() const;
+
   /// Transform the requirements into a form where implicit Copyable and
   /// Escapable conformances are omitted, and their absence is explicitly
   /// noted.

include/swift/AST/LocalArchetypeRequirementCollector.h

@@ -0,0 +1,68 @@
+//===--- LocalArchetypeRequirementCollector.h -------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file has utility code for extending a generic signature with opened
+// existentials and shape classes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_LOCAL_ARCHETYPE_REQUIREMENT_COLLECTOR_H
+#define SWIFT_AST_LOCAL_ARCHETYPE_REQUIREMENT_COLLECTOR_H
+
+#include "swift/AST/ASTContext.h"
+#include "swift/AST/GenericSignature.h"
+#include "swift/AST/Requirement.h"
+#include "swift/AST/Types.h"
+
+namespace swift {
+
+struct LocalArchetypeRequirementCollector {
+  const ASTContext &Context;
+  GenericSignature OuterSig;
+  unsigned Depth;
+
+  /// The lists of new parameters and requirements to add to the signature.
+  SmallVector<GenericTypeParamType *, 2> Params;
+  SmallVector<Requirement, 2> Requirements;
+
+  LocalArchetypeRequirementCollector(const ASTContext &ctx, GenericSignature sig);
+
+  void addOpenedExistential(Type constraint);
+  void addOpenedElement(CanGenericTypeParamType shapeClass);
+
+  GenericTypeParamType *addParameter();
+};
+
+struct MapLocalArchetypesOutOfContext {
+  GenericSignature baseGenericSig;
+  ArrayRef<GenericEnvironment *> capturedEnvs;
+
+  MapLocalArchetypesOutOfContext(GenericSignature baseGenericSig,
+                                 ArrayRef<GenericEnvironment *> capturedEnvs)
+      : baseGenericSig(baseGenericSig), capturedEnvs(capturedEnvs) {}
+
+  Type operator()(SubstitutableType *type) const;
+};
+
+GenericSignature buildGenericSignatureWithCapturedEnvironments(
+    ASTContext &ctx,
+    GenericSignature sig,
+    ArrayRef<GenericEnvironment *> capturedEnvs);
+
+SubstitutionMap buildSubstitutionMapWithCapturedEnvironments(
+    SubstitutionMap baseSubMap,
+    GenericSignature genericSigWithCaptures,
+    ArrayRef<GenericEnvironment *> capturedEnvs);
+
+}
+
+#endif // SWIFT_AST_LOCAL_ARCHETYPE_REQUIREMENT_COLLECTOR_H

lib/AST/ASTContext.cpp

@@ -36,6 +36,7 @@
 #include "swift/AST/IndexSubset.h"
 #include "swift/AST/KnownProtocols.h"
 #include "swift/AST/LazyResolver.h"
+#include "swift/AST/LocalArchetypeRequirementCollector.h"
 #include "swift/AST/MacroDiscriminatorContext.h"
 #include "swift/AST/ModuleDependencies.h"
 #include "swift/AST/ModuleLoader.h"
@@ -5884,11 +5885,8 @@ CanGenericSignature ASTContext::getSingleGenericParameterSignature() const {
 
 Type OpenedArchetypeType::getSelfInterfaceTypeFromContext(GenericSignature parentSig,
                                                           ASTContext &ctx) {
-  unsigned depth = 0;
-  if (!parentSig.getGenericParams().empty())
-    depth = parentSig.getGenericParams().back()->getDepth() + 1;
   return GenericTypeParamType::get(/*isParameterPack=*/ false,
-                                   /*depth=*/ depth, /*index=*/ 0,
+                                   parentSig.getNextDepth(), /*index=*/ 0,
                                    ctx);
 }
 
@@ -5900,33 +5898,25 @@ ASTContext::getOpenedExistentialSignature(Type type, GenericSignature parentSig)
     type = existential->getConstraintType();
 
   const CanType constraint = type->getCanonicalType();
-  assert(parentSig || !constraint->hasTypeParameter() &&
-         "Interface type here requires a parent signature");
 
   auto canParentSig = parentSig.getCanonicalSignature();
   auto key = std::make_pair(constraint, canParentSig.getPointer());
   auto found = getImpl().ExistentialSignatures.find(key);
   if (found != getImpl().ExistentialSignatures.end())
     return found->second;
 
-  auto genericParam = OpenedArchetypeType::getSelfInterfaceTypeFromContext(
-      canParentSig, *this)
-    ->castTo<GenericTypeParamType>();
-  Requirement requirement(RequirementKind::Conformance, genericParam,
-                          constraint);
+  LocalArchetypeRequirementCollector collector(*this, canParentSig);
+  collector.addOpenedExistential(type);
   auto genericSig = buildGenericSignature(
-      *this, canParentSig,
-      {genericParam}, {requirement},
-      /*allowInverses=*/true);
-
-  CanGenericSignature canGenericSig(genericSig);
+      *this, collector.OuterSig, collector.Params, collector.Requirements,
+      /*allowInverses=*/true).getCanonicalSignature();
 
   auto result = getImpl().ExistentialSignatures.insert(
-      std::make_pair(key, canGenericSig));
+      std::make_pair(key, genericSig));
   assert(result.second);
   (void) result;
 
-  return canGenericSig;
+  return genericSig;
 }
 
 CanGenericSignature
@@ -5938,100 +5928,12 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig,
   if (found != sigs.end())
     return found->second;
 
-  // This operation doesn't make sense if the input signature does not contain`
-  // any pack generic parameters.
-#ifndef NDEBUG
-  {
-    auto found = std::find_if(baseGenericSig.getGenericParams().begin(),
-                              baseGenericSig.getGenericParams().end(),
-                              [](GenericTypeParamType *paramType) {
-                                return paramType->isParameterPack();
-                              });
-    assert(found != baseGenericSig.getGenericParams().end());
-  }
-#endif
-
-  // The pack element signature includes all type parameters and requirements
-  // from the outer context, plus a new set of type parameters representing
-  // open pack elements and their corresponding element requirements.
-
-  llvm::SmallMapVector<GenericTypeParamType *,
-                       GenericTypeParamType *, 2> packElementParams;
-  SmallVector<GenericTypeParamType *, 2> genericParams(
-      baseGenericSig.getGenericParams().begin(), baseGenericSig.getGenericParams().end());
-  SmallVector<Requirement, 2> requirements;
-
-  auto packElementDepth =
-      baseGenericSig.getInnermostGenericParams().front()->getDepth() + 1;
-
-  for (auto paramType : baseGenericSig.getGenericParams()) {
-    if (!paramType->isParameterPack())
-      continue;
-
-    // Only include opened element parameters for packs in the given
-    // shape equivalence class.
-    if (!baseGenericSig->haveSameShape(paramType, shapeClass))
-      continue;
-
-    auto *elementParam = GenericTypeParamType::get(/*isParameterPack*/false,
-                                                   packElementDepth,
-                                                   packElementParams.size(),
-                                                   *this);
-    genericParams.push_back(elementParam);
-    packElementParams[paramType] = elementParam;
-  }
-
-  auto eraseParameterPackRec = [&](Type type) -> Type {
-    return type.transformTypeParameterPacks(
-        [&](SubstitutableType *t) -> std::optional<Type> {
-          if (auto *paramType = dyn_cast<GenericTypeParamType>(t)) {
-            if (packElementParams.find(paramType) != packElementParams.end()) {
-              return Type(packElementParams[paramType]);
-            }
-
-            return Type(t);
-          }
-          return std::nullopt;
-        });
-  };
-
-  for (auto requirement : baseGenericSig.getRequirements()) {
-    requirements.push_back(requirement);
-
-    // If this requirement contains parameter packs, create a new requirement
-    // for the corresponding pack element.
-    switch (requirement.getKind()) {
-    case RequirementKind::SameShape:
-      // Drop same-shape requirements from the element signature.
-      break;
-    case RequirementKind::Conformance:
-    case RequirementKind::Superclass:
-    case RequirementKind::SameType: {
-      auto firstType = eraseParameterPackRec(requirement.getFirstType());
-      auto secondType = eraseParameterPackRec(requirement.getSecondType());
-      if (firstType->isEqual(requirement.getFirstType()) &&
-          secondType->isEqual(requirement.getSecondType()))
-        break;
-
-      requirements.emplace_back(requirement.getKind(),
-                                firstType, secondType);
-      break;
-    }
-    case RequirementKind::Layout: {
-      auto firstType = eraseParameterPackRec(requirement.getFirstType());
-      if (firstType->isEqual(requirement.getFirstType()))
-        break;
-
-      requirements.emplace_back(requirement.getKind(), firstType,
-                                requirement.getLayoutConstraint());
-      break;
-    }
-    }
-  }
-
+  LocalArchetypeRequirementCollector collector(*this, baseGenericSig);
+  collector.addOpenedElement(shapeClass);
   auto elementSig = buildGenericSignature(
-      *this, GenericSignature(), genericParams, requirements,
+      *this, collector.OuterSig, collector.Params, collector.Requirements,
       /*allowInverses=*/false).getCanonicalSignature();
+
   sigs[key] = elementSig;
   return elementSig;
 }

lib/AST/ASTMangler.cpp

@@ -2748,7 +2748,7 @@ void ASTMangler::appendContextualInverses(const GenericTypeDecl *contextDecl,
   parts.params = std::nullopt;
 
   // The depth of parameters for this extension is +1 of the extended signature.
-  parts.initialParamDepth = sig.getGenericParams().back()->getDepth() + 1;
+  parts.initialParamDepth = sig.getNextDepth();
 
   appendModule(module, alternateModuleName);
   appendGenericSignatureParts(sig, parts);
@@ -3391,7 +3391,7 @@ void ASTMangler::gatherGenericSignatureParts(GenericSignature sig,
   } else {
     inverseReqs.clear();
   }
-  base.setDepth(canSig.getGenericParams().back()->getDepth());
+  base.setDepth(canSig->getMaxDepth());
 
   unsigned &initialParamDepth = parts.initialParamDepth;
   auto &genericParams = parts.params;
@@ -3412,7 +3412,7 @@ void ASTMangler::gatherGenericSignatureParts(GenericSignature sig,
 
   // The signature depth starts above the depth of the context signature.
   if (!contextSig.getGenericParams().empty()) {
-    initialParamDepth = contextSig.getGenericParams().back()->getDepth() + 1;
+    initialParamDepth = contextSig.getNextDepth();
   }
 
   // If both signatures have exactly the same requirements, ignoring
@@ -4760,11 +4760,10 @@ static std::optional<unsigned> getEnclosingTypeGenericDepth(const Decl *decl) {
   if (!typeDecl)
     return std::nullopt;
 
-  auto genericParams = typeDecl->getGenericParams();
-  if (!genericParams)
+  if (!typeDecl->isGeneric())
     return std::nullopt;
 
-  return genericParams->getParams().back()->getDepth();
+  return typeDecl->getGenericSignature()->getMaxDepth();
 }
 
 ASTMangler::BaseEntitySignature::BaseEntitySignature(const Decl *decl)

lib/AST/ASTPrinter.cpp

@@ -1807,10 +1807,7 @@ void PrintAST::printSingleDepthOfGenericSignature(
       auto *DC = Current->getInnermostDeclContext()->getInnermostTypeContext();
       M = DC->getParentModule();
       subMap = CurrentType->getContextSubstitutionMap(M, DC);
-      if (!subMap.empty()) {
-        typeContextDepth = subMap.getGenericSignature()
-            .getGenericParams().back()->getDepth() + 1;
-      }
+      typeContextDepth = subMap.getGenericSignature().getNextDepth();
     }
   }
 
@@ -7080,7 +7077,7 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
 
     // The element archetypes are at a depth one past the max depth
     // of the base signature.
-    unsigned elementDepth = params.back()->getDepth() + 1;
+    unsigned elementDepth = sig.getNextDepth();
 
     // Transform the archetype's interface type to be based on the
     // corresponding non-canonical type parameter.

lib/AST/CMakeLists.txt

@@ -62,6 +62,7 @@ add_swift_host_library(swiftAST STATIC
   InlinableText.cpp
   LayoutConstraint.cpp
   LifetimeDependence.cpp
+  LocalArchetypeRequirementCollector.cpp
   Module.cpp
   ModuleDependencies.cpp
   ModuleLoader.cpp

lib/AST/GenericSignature.cpp

@@ -79,7 +79,7 @@ ArrayRef<GenericTypeParamType *>
 GenericSignatureImpl::getInnermostGenericParams() const {
   const auto params = getGenericParams();
 
-  const unsigned maxDepth = params.back()->getDepth();
+  const unsigned maxDepth = getMaxDepth();
   if (params.front()->getDepth() == maxDepth)
     return params;
 
@@ -95,6 +95,16 @@ GenericSignatureImpl::getInnermostGenericParams() const {
   return params.slice(sliceCount);
 }
 
+unsigned GenericSignatureImpl::getMaxDepth() const {
+  return getGenericParams().back()->getDepth();
+}
+
+unsigned GenericSignature::getNextDepth() const {
+  if (!getPointer())
+    return 0;
+  return getPointer()->getMaxDepth() + 1;
+}
+
 void GenericSignatureImpl::forEachParam(
     llvm::function_ref<void(GenericTypeParamType *, bool)> callback) const {
   // Figure out which generic parameters are concrete or same-typed to another