[ConstraintSystem] implement implicit pack materialization for abstract tuples instead of explicit '.element'

include/swift/AST/KnownIdentifiers.def

@@ -325,7 +325,6 @@ IDENTIFIER(size)
 IDENTIFIER(speed)
 IDENTIFIER(unchecked)
 IDENTIFIER(unsafe)
-IDENTIFIER(element)
 
 // The singleton instance of TupleTypeDecl in the Builtin module
 IDENTIFIER(TheTupleType)

include/swift/Sema/Constraint.h

@@ -234,6 +234,9 @@ enum class ConstraintKind : char {
   ExplicitGenericArguments,
   /// Both (first and second) pack types should have the same reduced shape.
   SameShape,
+  /// The first type is a tuple containing a single unlabeled element that is a
+  /// pack expansion. The second type is that pack expansion.
+  MaterializePackExpansion,
 };
 
 /// Classification of the different kinds of constraints.
@@ -703,6 +706,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::UnresolvedMemberChainBase:
     case ConstraintKind::PackElementOf:
     case ConstraintKind::SameShape:
+    case ConstraintKind::MaterializePackExpansion:
       return ConstraintClassification::Relational;
 
     case ConstraintKind::ValueMember:

include/swift/Sema/ConstraintSystem.h

@@ -1475,6 +1475,10 @@ class Solution {
   llvm::DenseMap<ConstraintLocator *, std::pair<UUID, Type>>
       PackExpansionEnvironments;
 
+  /// The pack expansion environment that can open a given pack element.
+  llvm::SmallMapVector<PackElementExpr *, PackExpansionExpr *, 2>
+      PackEnvironments;
+
   /// The locators of \c Defaultable constraints whose defaults were used.
   llvm::SmallPtrSet<ConstraintLocator *, 2> DefaultedConstraints;
 
@@ -2251,6 +2255,9 @@ class ConstraintSystem {
   llvm::SmallMapVector<ConstraintLocator *, std::pair<UUID, Type>, 4>
       PackExpansionEnvironments;
 
+  llvm::SmallMapVector<PackElementExpr *, PackExpansionExpr *, 2>
+      PackEnvironments;
+
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
       resultBuilderTransformed;
@@ -2737,6 +2744,9 @@ class ConstraintSystem {
     /// The length of \c PackExpansionEnvironments.
     unsigned numPackExpansionEnvironments;
 
+    /// The length of \c PackEnvironments.
+    unsigned numPackEnvironments;
+
     /// The length of \c DefaultedConstraints.
     unsigned numDefaultedConstraints;
 
@@ -3232,6 +3242,13 @@ class ConstraintSystem {
   GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
                                                 CanType shapeClass);
 
+  /// Get the opened element generic environment for the given pack element.
+  PackExpansionExpr *getPackEnvironment(PackElementExpr *packElement) const;
+
+  /// Associate an opened element generic environment to a pack element.
+  void addPackEnvironment(PackElementExpr *packElement,
+                          PackExpansionExpr *packExpansion);
+
   /// Retrieve the constraint locator for the given anchor and
   /// path, uniqued and automatically infer the summary flags
   ConstraintLocator *
@@ -4921,6 +4938,13 @@ class ConstraintSystem {
                                            TypeMatchOptions flags,
                                            ConstraintLocatorBuilder locator);
 
+  /// Remove the tuple wrapping of left-hand type if it contains only a single
+  /// unlabeled element that is a pack expansion.
+  SolutionKind
+  simplifyMaterializePackExpansionConstraint(Type type1, Type type2,
+                                             TypeMatchOptions flags,
+                                             ConstraintLocatorBuilder locator);
+
 public: // FIXME: Public for use by static functions.
   /// Simplify a conversion constraint with a fix applied to it.
   SolutionKind simplifyFixConstraint(ConstraintFix *fix, Type type1, Type type2,
@@ -5463,7 +5487,8 @@ class OpenPackElementType {
     }
 
     cs.addConstraint(ConstraintKind::PackElementOf, elementType,
-                     packType, cs.getConstraintLocator(elementEnv));
+                     packType->getRValueType(),
+                     cs.getConstraintLocator(elementEnv));
     return elementType;
   }
 };
@@ -6146,6 +6171,9 @@ void dumpAnchor(ASTNode anchor, SourceManager *SM, raw_ostream &out);
 
 bool isSingleUnlabeledPackExpansionTuple(Type type);
 
+/// \returns null if \c type is not a single unlabeled pack expansion tuple.
+Type getPatternTypeOfSingleUnlabeledPackExpansionTuple(Type type);
+
 } // end namespace constraints
 
 template<typename ...Args>

lib/Sema/CSApply.cpp

@@ -3852,9 +3852,8 @@ namespace {
         auto packRefType = cs.getType(packRefExpr);
         if (auto *tuple = packRefType->getRValueType()->getAs<TupleType>();
             tuple && tuple->isSingleUnlabeledPackExpansion()) {
-          auto *expansion =
-              tuple->getElementType(0)->castTo<PackExpansionType>();
-          auto patternType = expansion->getPatternType();
+          auto patternType =
+              getPatternTypeOfSingleUnlabeledPackExpansionTuple(tuple);
           auto *materializedPackExpr = MaterializePackExpr::create(
               cs.getASTContext(), packRefExpr, packRefExpr->getLoc(),
               patternType, /*implicit*/ true);

lib/Sema/CSBindings.cpp

@@ -1483,6 +1483,7 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::ExplicitGenericArguments:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::SameShape:
+  case ConstraintKind::MaterializePackExpansion:
     // Constraints from which we can't do anything.
     break;
 

lib/Sema/CSDiagnostics.cpp

@@ -6170,8 +6170,10 @@ bool InvalidPackElement::diagnoseAsError() {
 }
 
 bool InvalidPackReference::diagnoseAsError() {
-  emitDiagnostic(diag::pack_reference_outside_expansion,
-                 packType);
+  auto patternType =
+      getPatternTypeOfSingleUnlabeledPackExpansionTuple(packType);
+  auto diagnosisType = patternType ? patternType : packType;
+  emitDiagnostic(diag::pack_reference_outside_expansion, diagnosisType);
   return true;
 }
 

lib/Sema/CSGen.cpp

@@ -1089,20 +1089,16 @@ namespace {
       return outputTy;
     }
 
-    Type openPackElement(Type packType, ConstraintLocator *locator) {
-      // If 'each t' is written outside of a pack expansion expression, allow the
-      // type to bind to a hole. The invalid pack reference will be diagnosed when
-      // attempting to bind the type variable for the underlying pack reference to
-      // a pack type without TVO_CanBindToPack.
-      if (PackElementEnvironments.empty()) {
+    Type openPackElement(Type packType, ConstraintLocator *locator,
+                         PackExpansionExpr *packElementEnvironment) {
+      if (!packElementEnvironment) {
         return CS.createTypeVariable(locator,
                                      TVO_CanBindToHole | TVO_CanBindToNoEscape);
       }
 
       // The type of a PackElementExpr is the opened pack element archetype
       // of the pack reference.
-      OpenPackElementType openPackElement(CS, locator,
-                                          PackElementEnvironments.back());
+      OpenPackElementType openPackElement(CS, locator, packElementEnvironment);
       return openPackElement(packType, /*packRepr*/ nullptr);
     }
 
@@ -1124,6 +1120,26 @@ namespace {
 
     void addPackElementEnvironment(PackExpansionExpr *expr) {
       PackElementEnvironments.push_back(expr);
+
+      SmallVector<ASTNode, 2> expandedPacks;
+      collectExpandedPacks(expr, expandedPacks);
+      for (auto pack : expandedPacks) {
+        if (auto *elementExpr = getAsExpr<PackElementExpr>(pack)) {
+          CS.addPackEnvironment(elementExpr, expr);
+        }
+      }
+
+      auto *patternLoc = CS.getConstraintLocator(
+          expr, ConstraintLocator::PackExpansionPattern);
+      auto patternType = CS.createTypeVariable(
+          patternLoc,
+          TVO_CanBindToPack | TVO_CanBindToNoEscape | TVO_CanBindToHole);
+      auto *shapeLoc =
+          CS.getConstraintLocator(expr, ConstraintLocator::PackShape);
+      auto *shapeTypeVar = CS.createTypeVariable(
+          shapeLoc, TVO_CanBindToPack | TVO_CanBindToHole);
+      auto expansionType = PackExpansionType::get(patternType, shapeTypeVar);
+      CS.setType(expr, expansionType);
     }
 
     virtual Type visitErrorExpr(ErrorExpr *E) {
@@ -1384,6 +1400,17 @@ namespace {
       return BoundGenericStructType::get(regexDecl, Type(), {matchType});
     }
 
+    PackExpansionExpr *getParentPackExpansionExpr(Expr *E) const {
+      auto *current = E;
+      while (auto *parent = CS.getParentExpr(current)) {
+        if (auto *expansion = dyn_cast<PackExpansionExpr>(parent)) {
+          return expansion;
+        }
+        current = parent;
+      }
+      return nullptr;
+    }
+
     Type visitDeclRefExpr(DeclRefExpr *E) {
       auto locator = CS.getConstraintLocator(E);
 
@@ -1426,13 +1453,15 @@ namespace {
 
           // value packs cannot be referenced without `each` immediately
           // preceding them.
-          if (auto *expansion = knownType->getAs<PackExpansionType>()) {
-            if (!PackElementEnvironments.empty() &&
+          if (auto *expansionType = knownType->getAs<PackExpansionType>()) {
+            if (auto *parentExpansionExpr = getParentPackExpansionExpr(E);
+                parentExpansionExpr &&
                 !isExpr<PackElementExpr>(CS.getParentExpr(E))) {
-              auto packType = expansion->getPatternType();
+              auto packType = expansionType->getPatternType();
               (void)CS.recordFix(
                   IgnoreMissingEachKeyword::create(CS, packType, locator));
-              auto eltType = openPackElement(packType, locator);
+              auto eltType =
+                  openPackElement(packType, locator, parentExpansionExpr);
               CS.setType(E, eltType);
               return eltType;
             }
@@ -3033,21 +3062,11 @@ namespace {
       assert(PackElementEnvironments.back() == expr);
       PackElementEnvironments.pop_back();
 
-      auto *patternLoc =
-          CS.getConstraintLocator(expr, ConstraintLocator::PackExpansionPattern);
-      auto patternTy = CS.createTypeVariable(patternLoc,
-                                             TVO_CanBindToPack |
-                                             TVO_CanBindToNoEscape |
-                                             TVO_CanBindToHole);
+      auto expansionType = CS.getType(expr)->castTo<PackExpansionType>();
       auto elementResultType = CS.getType(expr->getPatternExpr());
       CS.addConstraint(ConstraintKind::PackElementOf, elementResultType,
-                       patternTy, CS.getConstraintLocator(expr));
-
-      auto *shapeLoc =
-          CS.getConstraintLocator(expr, ConstraintLocator::PackShape);
-      auto *shapeTypeVar = CS.createTypeVariable(shapeLoc,
-                                                 TVO_CanBindToPack |
-                                                 TVO_CanBindToHole);
+                       expansionType->getPatternType(),
+                       CS.getConstraintLocator(expr));
 
       // Generate ShapeOf constraints between all packs expanded by this
       // pack expansion expression through the shape type variable.
@@ -3061,9 +3080,14 @@ namespace {
 
       for (auto pack : expandedPacks) {
         Type packType;
-        if (auto *elementExpr = getAsExpr<PackElementExpr>(pack)) {
-          packType = CS.getType(elementExpr->getPackRefExpr());
-        } else if (auto *elementType = getAsTypeRepr<PackElementTypeRepr>(pack)) {
+        /// Skipping over pack elements because the relationship to its
+        /// environment is now established during \c addPackElementEnvironment
+        /// upon visiting its pack expansion and the Shape constraint added
+        /// upon visiting the pack element.
+        if (isExpr<PackElementExpr>(pack)) {
+          continue;
+        } else if (auto *elementType =
+                       getAsTypeRepr<PackElementTypeRepr>(pack)) {
           // OpenPackElementType sets types for 'each T' type reprs in
           // expressions. Some invalid code won't make it there, and
           // the constraint system won't have recorded a type.
@@ -3075,26 +3099,47 @@ namespace {
           llvm_unreachable("unsupported pack reference ASTNode");
         }
 
+        auto *elementShape = CS.createTypeVariable(
+            CS.getConstraintLocator(pack, ConstraintLocator::PackShape),
+            TVO_CanBindToPack);
+        CS.addConstraint(
+            ConstraintKind::ShapeOf, elementShape, packType,
+            CS.getConstraintLocator(pack, ConstraintLocator::PackShape));
         CS.addConstraint(
-            ConstraintKind::ShapeOf, shapeTypeVar, packType,
+            ConstraintKind::Equal, elementShape, expansionType->getCountType(),
             CS.getConstraintLocator(expr, ConstraintLocator::PackShape));
       }
 
-      return PackExpansionType::get(patternTy, shapeTypeVar);
+      return expansionType;
     }
 
     Type visitPackElementExpr(PackElementExpr *expr) {
       auto packType = CS.getType(expr->getPackRefExpr());
-
-      if (isSingleUnlabeledPackExpansionTuple(packType)) {
-        packType =
-            addMemberRefConstraints(expr, expr->getPackRefExpr(),
-                                    DeclNameRef(CS.getASTContext().Id_element),
-                                    FunctionRefKind::Unapplied, {});
-        CS.setType(expr->getPackRefExpr(), packType);
+      auto *packEnvironment = CS.getPackEnvironment(expr);
+      auto elementType = openPackElement(
+          packType, CS.getConstraintLocator(expr), packEnvironment);
+      if (packEnvironment) {
+        auto expansionType =
+            CS.getType(packEnvironment)->castTo<PackExpansionType>();
+        CS.addConstraint(ConstraintKind::ShapeOf, expansionType->getCountType(),
+                         elementType,
+                         CS.getConstraintLocator(packEnvironment,
+                                                 ConstraintLocator::PackShape));
+        auto *elementShape = CS.createTypeVariable(
+            CS.getConstraintLocator(expr, ConstraintLocator::PackShape),
+            TVO_CanBindToPack);
+        CS.addConstraint(
+            ConstraintKind::ShapeOf, elementShape, elementType,
+            CS.getConstraintLocator(expr, ConstraintLocator::PackShape));
+        CS.addConstraint(
+            ConstraintKind::Equal, elementShape, expansionType->getCountType(),
+            CS.getConstraintLocator(expr, ConstraintLocator::PackShape));
+      } else {
+        CS.recordFix(AllowInvalidPackReference::create(
+            CS, packType, CS.getConstraintLocator(expr->getPackRefExpr())));
       }
 
-      return openPackElement(packType, CS.getConstraintLocator(expr));
+      return elementType;
     }
 
     Type visitMaterializePackExpr(MaterializePackExpr *expr) {