Merge pull request #64119 from hborla/pack-expansion-shape-crash

[Diagnostics] Diagnose pack element expressions containing a non-pack a subexpression.

Author: hborla

include/swift/Sema/CSFix.h

@@ -425,6 +425,9 @@ enum class FixKind : uint8_t {
 
   /// Produce an error about a type that must be Copyable
   MustBeCopyable,
+
+  /// Allow 'each' applied to a non-pack type.
+  AllowInvalidPackElement,
 };
 
 class ConstraintFix {
@@ -2061,6 +2064,30 @@ class MustBeCopyable final : public ConstraintFix {
   }
 };
 
+class AllowInvalidPackElement final : public ConstraintFix {
+  Type packElementType;
+
+  AllowInvalidPackElement(ConstraintSystem &cs, Type packElementType,
+                          ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::AllowInvalidPackElement, locator),
+        packElementType(packElementType) {}
+
+public:
+  std::string getName() const override {
+    return "allow concrete pack element";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  static AllowInvalidPackElement *create(ConstraintSystem &cs,
+                                         Type packElementType,
+                                         ConstraintLocator *locator);
+
+  static bool classof(ConstraintFix const* fix) {
+    return fix->getKind() == FixKind::AllowInvalidPackElement;
+  }
+};
+
 class CollectionElementContextualMismatch final
     : public ContextualMismatch,
       private llvm::TrailingObjects<CollectionElementContextualMismatch,

lib/AST/ASTContext.cpp

@@ -3247,17 +3247,7 @@ PackExpansionType::PackExpansionType(Type patternType, Type countType,
                                      RecursiveTypeProperties properties,
                                      const ASTContext *canCtx)
   : TypeBase(TypeKind::PackExpansion, canCtx, properties),
-    patternType(patternType), countType(countType) {
-
-  // TODO: it would be nice if the solver didn't stick PlaceholderTypes and
-  // UnresolvedTypes in here.
-  assert(countType->is<PackType>() ||
-         countType->is<TypeVariableType>() ||
-         countType->is<PackArchetypeType>() ||
-         countType->is<PlaceholderType>() ||
-         countType->is<UnresolvedType>() ||
-         countType->castTo<GenericTypeParamType>()->isParameterPack());
-}
+    patternType(patternType), countType(countType) {}
 
 CanPackExpansionType
 CanPackExpansionType::get(CanType patternType, CanType countType) {

lib/AST/ASTVerifier.cpp

@@ -603,6 +603,18 @@ class Verifier : public ASTWalker {
         abort();
       }
 
+      // Check for invalid pack expansion shape types.
+      if (auto *expansion = type->getAs<PackExpansionType>()) {
+        auto countType = expansion->getCountType();
+        if (!(countType->is<PackType>() ||
+              countType->is<PackArchetypeType>() ||
+              (countType->is<GenericTypeParamType>() &&
+               countType->castTo<GenericTypeParamType>()->isParameterPack()))) {
+          Out << "non-pack shape type" << countType->getString() << "\n";
+          abort();
+        }
+      }
+
       if (!type->hasArchetype())
         return;
 

lib/Sema/CSBindings.cpp

@@ -1494,6 +1494,12 @@ void PotentialBindings::infer(Constraint *constraint) {
       packType = packType->mapTypeOutOfContext();
       auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator(),
                                                       shapeClass);
+
+      // Without an opened element environment, we cannot derive the
+      // element binding.
+      if (!elementEnv)
+        break;
+
       auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
       addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});
 

lib/Sema/CSDiagnostics.cpp

@@ -5996,6 +5996,11 @@ bool NotCopyableFailure::diagnoseAsError() {
   return true;
 }
 
+bool InvalidPackElement::diagnoseAsError() {
+  emitDiagnostic(diag::each_non_pack, packElementType);
+  return true;
+}
+
 bool CollectionElementContextualFailure::diagnoseAsError() {
   auto anchor = getRawAnchor();
   auto *locator = getLocator();

lib/Sema/CSDiagnostics.h

@@ -1838,6 +1838,19 @@ class NotCopyableFailure final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+/// Diagnose \c each applied to an expression that is not a pack type.
+class InvalidPackElement final : public FailureDiagnostic {
+  Type packElementType;
+
+public:
+  InvalidPackElement(const Solution &solution, Type packElementType,
+                     ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator),
+        packElementType(packElementType) {}
+
+  bool diagnoseAsError() override;
+};
+
 /// Diagnose a contextual mismatch between expected collection element type
 /// and the one provided (e.g. source of the assignment or argument to a call)
 /// e.g.:

lib/Sema/CSFix.cpp

@@ -1366,6 +1366,20 @@ bool MustBeCopyable::diagnoseForAmbiguity(CommonFixesArray commonFixes) const {
   return diagnose(*commonFixes.front().first);
 }
 
+bool AllowInvalidPackElement::diagnose(const Solution &solution,
+                                       bool asNote) const {
+  InvalidPackElement failure(solution, packElementType, getLocator());
+  return failure.diagnose(asNote);
+}
+
+AllowInvalidPackElement *
+AllowInvalidPackElement::create(ConstraintSystem &cs,
+                                Type packElementType,
+                                ConstraintLocator *locator) {
+  return new (cs.getAllocator())
+      AllowInvalidPackElement(cs, packElementType, locator);
+}
+
 bool CollectionElementContextualMismatch::diagnose(const Solution &solution,
                                                    bool asNote) const {
   CollectionElementContextualFailure failure(

lib/Sema/CSSimplify.cpp

@@ -8818,6 +8818,13 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
     return SolutionKind::Solved;
   }
 
+  if (shouldAttemptFixes()) {
+    auto *loc = getConstraintLocator(locator);
+    auto *fix = AllowInvalidPackElement::create(*this, packType, loc);
+    if (!recordFix(fix))
+      return SolutionKind::Solved;
+  }
+
   return SolutionKind::Error;
 }
 
@@ -14029,6 +14036,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::NotCompileTimeConst:
   case FixKind::RenameConflictingPatternVariables:
   case FixKind::MustBeCopyable:
+  case FixKind::AllowInvalidPackElement:
   case FixKind::MacroMissingPound:
   case FixKind::AllowGlobalActorMismatch: {
     return recordFix(fix) ? SolutionKind::Error : SolutionKind::Solved;

lib/Sema/ConstraintSystem.cpp

@@ -661,6 +661,9 @@ ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
   if (result == PackExpansionEnvironments.end())
     return nullptr;
 
+  if (!shapeClass->is<PackArchetypeType>())
+    return nullptr;
+
   auto uuid = result->second;
   auto &ctx = getASTContext();
   auto elementSig = ctx.getOpenedElementSignature(

test/Constraints/pack-expansion-expressions.swift

@@ -120,4 +120,8 @@ func generateTuple<each T : Generatable>() -> (repeat each T) {
 
 func packElementInvalidBinding<each T>(_ arg: repeat each T) {
   _ = (repeat print(each arg))
+
+  let x = 1
+  repeat print(each x)
+  // expected-error@-1 {{'each' cannot be applied to non-pack type 'Int'}}
 }