Merge pull request #64166 from hborla/parameter-pack-substitution

hborla · web-flow · commit 75c1b6e199fc · 2023-03-07T06:39:50.000-08:00
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -2733,6 +2733,9 @@ ERROR(requires_not_suitable_archetype,none,
 ERROR(invalid_shape_requirement,none,
       "invalid same-shape requirement between %0 and %1",
       (Type, Type))
+ERROR(unsupported_same_element,none,
+      "same-element requirements are not yet supported",
+      ())
 
 ERROR(requires_generic_params_made_equal,none,
       "same-type requirement makes generic parameters %0 and %1 equivalent",
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -1487,7 +1487,8 @@ class Solution {
 
   /// The pack expansion environment that can open pack elements for
   /// a given locator.
-  llvm::DenseMap<ConstraintLocator *, UUID> PackExpansionEnvironments;
+  llvm::DenseMap<ConstraintLocator *, std::pair<UUID, Type>>
+      PackExpansionEnvironments;
 
   /// The locators of \c Defaultable constraints whose defaults were used.
   llvm::SmallPtrSet<ConstraintLocator *, 2> DefaultedConstraints;
@@ -2228,7 +2229,8 @@ class ConstraintSystem {
   llvm::SmallMapVector<ConstraintLocator *, OpenedArchetypeType *, 4>
       OpenedExistentialTypes;
 
-  llvm::SmallMapVector<ConstraintLocator *, UUID, 4> PackExpansionEnvironments;
+  llvm::SmallMapVector<ConstraintLocator *, std::pair<UUID, Type>, 4>
+      PackExpansionEnvironments;
 
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
@@ -3191,9 +3193,6 @@ class ConstraintSystem {
   std::pair<Type, OpenedArchetypeType *> openExistentialType(
       Type type, ConstraintLocator *locator);
 
-  /// Add the given pack expansion as an opened pack element environment.
-  void addPackElementEnvironment(PackExpansionExpr *expr);
-
   /// Get the opened element generic environment for the given locator.
   GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
                                                 CanType shapeClass);
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -5147,6 +5147,7 @@ GenericEnvironment::forOpenedElement(GenericSignature signature,
   if (found != openedElementEnvironments.end()) {
     auto *existingEnv = found->second;
     assert(existingEnv->getGenericSignature().getPointer() == signature.getPointer());
+    assert(existingEnv->getOpenedElementShapeClass()->isEqual(shapeClass));
     assert(existingEnv->getOpenedElementUUID() == uuid);
 
     return existingEnv;
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -4670,8 +4670,12 @@ static Type computeNominalType(NominalTypeDecl *decl, DeclTypeKind kind) {
       // the generic parameter list directly instead of looking
       // at the signature.
       SmallVector<Type, 4> args;
-      for (auto param : decl->getGenericParams()->getParams())
-        args.push_back(param->getDeclaredInterfaceType());
+      for (auto param : decl->getGenericParams()->getParams()) {
+        auto argTy = param->getDeclaredInterfaceType();
+        if (param->isParameterPack())
+          argTy = PackType::getSingletonPackExpansion(argTy);
+        args.push_back(argTy);
+      }
 
       return BoundGenericType::get(decl, ParentTy, args);
     }
diff --git a/lib/AST/GenericEnvironment.cpp b/lib/AST/GenericEnvironment.cpp
@@ -681,10 +681,37 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
   }, LookUpConformanceInSignature(sig.getPointer()));
 }
 
+namespace {
+/// A function suitable for use as a \c TypeSubstitutionFn that produces
+/// correct forwarding substitutions for a generic environment.
+///
+/// This differs from QueryInterfaceTypeSubstitutions only in that it
+/// always produces PackTypes for pack parameters.
+class BuildForwardingSubstitutions {
+  QueryInterfaceTypeSubstitutions Query;
+
+public:
+  BuildForwardingSubstitutions(const GenericEnvironment *self)
+    : Query(self) { }
+
+  Type operator()(SubstitutableType *type) const;
+};
+} // end anonymous namespace
+
+Type BuildForwardingSubstitutions::operator()(SubstitutableType *type) const {
+  if (auto resultType = Query(type)) {
+    auto param = type->castTo<GenericTypeParamType>();
+    if (!param->isParameterPack())
+      return resultType;
+    return PackType::getSingletonPackExpansion(resultType);
+  }
+  return Type();
+}
+
 SubstitutionMap GenericEnvironment::getForwardingSubstitutionMap() const {
   auto genericSig = getGenericSignature();
   return SubstitutionMap::get(genericSig,
-                              QueryInterfaceTypeSubstitutions(this),
+                              BuildForwardingSubstitutions(this),
                               MakeAbstractConformanceForGenericType());
 }
 
diff --git a/lib/AST/GenericSignature.cpp b/lib/AST/GenericSignature.cpp
@@ -579,10 +579,13 @@ unsigned GenericParamKey::findIndexIn(
 
 SubstitutionMap GenericSignatureImpl::getIdentitySubstitutionMap() const {
   return SubstitutionMap::get(const_cast<GenericSignatureImpl *>(this),
-                              [](SubstitutableType *t) -> Type {
-                                return Type(cast<GenericTypeParamType>(t));
-                              },
-                              MakeAbstractConformanceForGenericType());
+    [](SubstitutableType *t) -> Type {
+      auto param = cast<GenericTypeParamType>(t);
+      if (!param->isParameterPack())
+        return param;
+      return PackType::getSingletonPackExpansion(param);
+    },
+    MakeAbstractConformanceForGenericType());
 }
 
 GenericTypeParamType *GenericSignatureImpl::getSugaredType(
diff --git a/lib/AST/RequirementMachine/Diagnostics.cpp b/lib/AST/RequirementMachine/Diagnostics.cpp
@@ -209,6 +209,15 @@ bool swift::rewriting::diagnoseRequirementErrors(
 
       break;
     }
+
+    case RequirementError::Kind::UnsupportedSameElement: {
+      if (error.requirement.hasError())
+        break;
+
+      ctx.Diags.diagnose(loc, diag::unsupported_same_element);
+      diagnosedError = true;
+      break;
+    }
     }
   }
 
diff --git a/lib/AST/RequirementMachine/Diagnostics.h b/lib/AST/RequirementMachine/Diagnostics.h
@@ -42,6 +42,8 @@ struct RequirementError {
     RecursiveRequirement,
     /// A redundant requirement, e.g. T == T.
     RedundantRequirement,
+    /// A not-yet-supported same-element requirement, e.g. each T == Int.
+    UnsupportedSameElement,
   } kind;
 
   /// The invalid requirement.
@@ -100,6 +102,10 @@ struct RequirementError {
                                                   SourceLoc loc) {
     return {Kind::RecursiveRequirement, req, loc};
   }
+
+  static RequirementError forSameElement(Requirement req, SourceLoc loc) {
+    return {Kind::UnsupportedSameElement, req, loc};
+  }
 };
 
 /// Policy for the fixit that transforms 'T : S' where 'S' is not a protocol
diff --git a/lib/AST/RequirementMachine/RequirementLowering.cpp b/lib/AST/RequirementMachine/RequirementLowering.cpp
@@ -194,6 +194,15 @@ static void desugarSameTypeRequirement(Type lhs, Type rhs, SourceLoc loc,
 
     bool mismatch(TypeBase *firstType, TypeBase *secondType,
                   Type sugaredFirstType) {
+      // If one side is a parameter pack, this is a same-element requirement, which
+      // is not yet supported.
+      if (firstType->isParameterPack() != secondType->isParameterPack()) {
+        errors.push_back(RequirementError::forSameElement(
+            {RequirementKind::SameType, sugaredFirstType, secondType}, loc));
+        recordedErrors = true;
+        return true;
+      }
+
       if (firstType->isTypeParameter() && secondType->isTypeParameter()) {
         result.emplace_back(RequirementKind::SameType,
                             sugaredFirstType, secondType);
diff --git a/lib/AST/SubstitutionMap.cpp b/lib/AST/SubstitutionMap.cpp
@@ -226,6 +226,11 @@ SubstitutionMap SubstitutionMap::get(GenericSignature genericSig,
 
     // Record the replacement.
     Type replacement = Type(gp).subst(subs, lookupConformance);
+
+    assert((!replacement || replacement->hasError() ||
+            gp->isParameterPack() == replacement->is<PackType>()) &&
+           "replacement for pack parameter must be a pack type");
+
     replacementTypes.push_back(replacement);
   });
 
diff --git a/lib/AST/Type.cpp b/lib/AST/Type.cpp
@@ -4502,6 +4502,24 @@ operator()(CanType dependentType, Type conformingReplacementType,
 ProtocolConformanceRef MakeAbstractConformanceForGenericType::
 operator()(CanType dependentType, Type conformingReplacementType,
            ProtocolDecl *conformedProtocol) const {
+  // The places that use this can also produce conformance packs, generally
+  // just for singleton pack expansions.
+  if (auto conformingPack = conformingReplacementType->getAs<PackType>()) {
+    SmallVector<ProtocolConformanceRef, 4> conformances;
+    for (auto conformingPackElt : conformingPack->getElementTypes()) {
+      // Look through pack expansions; there's no equivalent conformance
+      // expansion right now.
+      auto expansion = conformingPackElt->getAs<PackExpansionType>();
+      if (expansion) conformingPackElt = expansion->getPatternType();
+
+      auto conformance =
+        (*this)(dependentType, conformingPackElt, conformedProtocol);
+      conformances.push_back(conformance);
+    }
+    return ProtocolConformanceRef(
+        PackConformance::get(conformingPack, conformedProtocol, conformances));
+  }
+
   assert((conformingReplacementType->is<ErrorType>() ||
           conformingReplacementType->is<SubstitutableType>() ||
           conformingReplacementType->is<DependentMemberType>() ||
@@ -4967,6 +4985,9 @@ TypeBase::getContextSubstitutions(const DeclContext *dc,
     else if (genericEnv)
       substTy = genericEnv->mapTypeIntoContext(gp);
 
+    if (gp->isParameterPack() && !substTy->hasError())
+      substTy = PackType::getSingletonPackExpansion(substTy);
+
     auto result = substitutions.insert(
       {gp->getCanonicalType()->castTo<GenericTypeParamType>(),
        substTy});
diff --git a/lib/IRGen/GenPack.cpp b/lib/IRGen/GenPack.cpp
@@ -263,6 +263,18 @@ static llvm::Value *bindWitnessTableAtIndex(IRGenFunction &IGF,
   return wtable;
 }
 
+/// Find the pack archetype for the given interface type in the given
+/// opened element context, which is known to be a forwarding context.
+static CanPackArchetypeType
+getMappedPackArchetypeType(const OpenedElementContext &context, CanType ty) {
+  auto packType = cast<PackType>(
+    context.environment->maybeApplyOuterContextSubstitutions(ty)
+        ->getCanonicalType());
+  auto archetype = getForwardedPackArchetypeType(packType);
+  assert(archetype);
+  return archetype;
+}
+
 static void bindElementSignatureRequirementsAtIndex(
     IRGenFunction &IGF, OpenedElementContext const &context, llvm::Value *index,
     DynamicMetadataRequest request) {
@@ -275,9 +287,7 @@ static void bindElementSignatureRequirementsAtIndex(
           break;
         case GenericRequirement::Kind::MetadataPack: {
           auto ty = requirement.getTypeParameter();
-          auto patternPackArchetype = cast<PackArchetypeType>(
-              context.environment->maybeApplyOuterContextSubstitutions(ty)
-                  ->getCanonicalType());
+          auto patternPackArchetype = getMappedPackArchetypeType(context, ty);
           auto response =
               IGF.emitTypeMetadataRef(patternPackArchetype, request);
           auto elementArchetype =
@@ -295,9 +305,7 @@ static void bindElementSignatureRequirementsAtIndex(
         case GenericRequirement::Kind::WitnessTablePack: {
           auto ty = requirement.getTypeParameter();
           auto proto = requirement.getProtocol();
-          auto patternPackArchetype = cast<PackArchetypeType>(
-              context.environment->maybeApplyOuterContextSubstitutions(ty)
-                  ->getCanonicalType());
+          auto patternPackArchetype = getMappedPackArchetypeType(context, ty);
           auto elementArchetype =
               context.environment
                   ->mapPackTypeIntoElementContext(
diff --git a/lib/SIL/IR/AbstractionPattern.cpp b/lib/SIL/IR/AbstractionPattern.cpp
@@ -1712,7 +1712,12 @@ class SubstFunctionTypePatternVisitor
     auto gp = GenericTypeParamType::get(isParameterPack, 0, paramIndex,
                                         TC.Context);
     substGenericParams.push_back(gp);
-    substReplacementTypes.push_back(substTy);
+    if (isParameterPack) {
+      substReplacementTypes.push_back(
+          PackType::getSingletonPackExpansion(substTy));
+    } else {
+      substReplacementTypes.push_back(substTy);
+    }
     
     if (auto layout = pattern.getLayoutConstraint()) {
       // Look at the layout constraint on this position in the abstraction pattern
diff --git a/lib/SILGen/SILGenConstructor.cpp b/lib/SILGen/SILGenConstructor.cpp
@@ -139,9 +139,13 @@ static SubstitutionMap getSubstitutionsForPropertyInitializer(
     // replacement types are the archetypes of the initializer itself.
     return SubstitutionMap::get(
       nominal->getGenericSignatureOfContext(),
-      [&](SubstitutableType *type) {
+      [&](SubstitutableType *type) -> Type {
         if (auto gp = type->getAs<GenericTypeParamType>()) {
-          return genericEnv->mapTypeIntoContext(gp);
+          auto archetype = genericEnv->mapTypeIntoContext(gp);
+          if (!gp->isParameterPack())
+            return archetype;
+
+          return PackType::getSingletonPackExpansion(archetype);
         }
 
         return Type(type);
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -95,8 +95,12 @@ Solution::computeSubstitutions(GenericSignature sig,
     return SubstitutionMap();
 
   TypeSubstitutionMap subs;
-  for (const auto &opened : openedTypes->second)
-    subs[opened.first] = getFixedType(opened.second);
+  for (const auto &opened : openedTypes->second) {
+    auto type = getFixedType(opened.second);
+    if (opened.first->isParameterPack() && !type->is<PackType>())
+      type = PackType::getSingletonPackExpansion(type);
+    subs[opened.first] = type;
+  }
 
   auto lookupConformanceFn =
       [&](CanType original, Type replacement,
diff --git a/lib/Sema/CSBindings.cpp b/lib/Sema/CSBindings.cpp
@@ -1501,6 +1501,7 @@ void PotentialBindings::infer(Constraint *constraint) {
         break;
 
       auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
+      assert(!elementType->is<PackType>());
       addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});
 
       break;
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -1101,7 +1101,6 @@ namespace {
     ConstraintSystem &getConstraintSystem() const { return CS; }
 
     void addPackElementEnvironment(PackExpansionExpr *expr) {
-      CS.addPackElementEnvironment(expr);
       PackElementEnvironments.push_back(expr);
     }
 
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -8843,6 +8843,19 @@ ConstraintSystem::simplifyBindTupleOfFunctionParamsConstraint(
   return SolutionKind::Solved;
 }
 
+static Type lookThroughSingletonPackExpansion(Type ty) {
+  if (auto pack = ty->getAs<PackType>()) {
+    if (pack->getNumElements() == 1) {
+      if (auto expansion = pack->getElementType(0)->getAs<PackExpansionType>()) {
+        auto countType = expansion->getCountType();
+        if (countType->isEqual(expansion->getPatternType()))
+          return countType;
+      }
+    }
+  }
+  return ty;
+}
+
 ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
                                                   TypeMatchOptions flags,
@@ -8862,6 +8875,11 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
     return SolutionKind::Solved;
   }
 
+  // FIXME: I'm not sure this is actually necessary; I may only be seeing
+  // this because of something I've screwed up in element generic
+  // environments.
+  elementType = lookThroughSingletonPackExpansion(elementType);
+
   // This constraint only exists to vend bindings.
   auto *packEnv = DC->getGenericEnvironmentOfContext();
 
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -649,29 +649,31 @@ std::pair<Type, OpenedArchetypeType *> ConstraintSystem::openExistentialType(
   return {result, opened};
 }
 
-void ConstraintSystem::addPackElementEnvironment(PackExpansionExpr *expr) {
-  auto *locator = getConstraintLocator(expr);
-  PackExpansionEnvironments[locator] = UUID::fromTime();
-}
-
 GenericEnvironment *
 ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
                                             CanType shapeClass) {
+  assert(locator->directlyAt<PackExpansionExpr>());
+
+  std::pair<UUID, Type> uuidAndShape;
   auto result = PackExpansionEnvironments.find(locator);
-  if (result == PackExpansionEnvironments.end())
-    return nullptr;
+  if (result == PackExpansionEnvironments.end()) {
+    uuidAndShape = std::make_pair(UUID::fromTime(), shapeClass);
+    PackExpansionEnvironments[locator] = uuidAndShape;
+  } else {
+    uuidAndShape = result->second;
+  }
 
-  if (!shapeClass->is<PackArchetypeType>())
+  if (!shapeClass->is<PackArchetypeType>() ||
+      !shapeClass->isEqual(uuidAndShape.second))
     return nullptr;
 
-  auto uuid = result->second;
   auto &ctx = getASTContext();
   auto elementSig = ctx.getOpenedElementSignature(
       DC->getGenericSignatureOfContext().getCanonicalSignature(), shapeClass);
   auto *contextEnv = DC->getGenericEnvironmentOfContext();
   auto contextSubs = contextEnv->getForwardingSubstitutionMap();
-  return GenericEnvironment::forOpenedElement(elementSig, uuid, shapeClass,
-                                              contextSubs);
+  return GenericEnvironment::forOpenedElement(elementSig, uuidAndShape.first,
+                                              shapeClass, contextSubs);
 }
 
 /// Extend the given depth map by adding depths for all of the subexpressions
diff --git a/lib/Sema/TypeCheckType.cpp b/lib/Sema/TypeCheckType.cpp
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
diff --git a/test/Constraints/pack-expansion-expressions.swift b/test/Constraints/pack-expansion-expressions.swift
diff --git a/test/SILGen/pack_expansion_type.swift b/test/SILGen/pack_expansion_type.swift
diff --git a/test/SILGen/variadic-generic-arguments.swift b/test/SILGen/variadic-generic-arguments.swift
diff --git a/test/SILGen/variadic-generic-results.swift b/test/SILGen/variadic-generic-results.swift
diff --git a/test/SILGen/variadic-generic-tuples.swift b/test/SILGen/variadic-generic-tuples.swift

Original file line number	Diff line number	Diff line change
`@@ -209,6 +209,15 @@ bool swift::rewriting::diagnoseRequirementErrors(`
`209`	`209`
`210`	`210`	`break;`
`211`	`211`	`}`
	`212`	`+`
	`213`	`+ case RequirementError::Kind::UnsupportedSameElement: {`
	`214`	`+ if (error.requirement.hasError())`
	`215`	`+ break;`
	`216`	`+`
	`217`	`+ ctx.Diags.diagnose(loc, diag::unsupported_same_element);`
	`218`	`+ diagnosedError = true;`
	`219`	`+ break;`
	`220`	`+ }`
`212`	`221`	`}`
`213`	`222`	`}`
`214`	`223`