Merge pull request #66214 from sophiapoirier/tuple-expansion-without-dot-element

[Variadic Generics] drop requirement of .element for tuple expansion

Author: sophiapoirier

include/swift/AST/Expr.h

@@ -3679,15 +3679,15 @@ class PackExpansionExpr final : public Expr {
 };
 
 /// An expression to materialize a pack from a tuple containing a pack
-/// expansion, spelled \c tuple.element.
+/// expansion.
 ///
 /// These nodes are created by CSApply and should only appear in a
 /// type-checked AST in the context of a \c PackExpansionExpr .
 class MaterializePackExpr final : public Expr {
   /// The expression from which to materialize a pack.
   Expr *FromExpr;
 
-  /// The source location of \c .element
+  /// The source location of (deprecated) \c .element
   SourceLoc ElementLoc;
 
   MaterializePackExpr(Expr *fromExpr, SourceLoc elementLoc,
@@ -3711,7 +3711,7 @@ class MaterializePackExpr final : public Expr {
   }
 
   SourceLoc getEndLoc() const {
-    return ElementLoc;
+    return ElementLoc.isInvalid() ? ElementLoc : FromExpr->getEndLoc();
   }
 
   static bool classof(const Expr *E) {

include/swift/AST/KnownIdentifiers.def

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

include/swift/Sema/ConstraintSystem.h

@@ -6100,6 +6100,8 @@ Type isPlaceholderVar(PatternBindingDecl *PB);
 /// Dump an anchor node for a constraint locator or contextual type.
 void dumpAnchor(ASTNode anchor, SourceManager *SM, raw_ostream &out);
 
+bool isSingleUnlabeledPackExpansionTuple(Type type);
+
 } // end namespace constraints
 
 template<typename ...Args>

include/swift/Sema/OverloadChoice.h

@@ -54,8 +54,7 @@ enum class OverloadChoiceKind : int {
   /// The overload choice selects a particular declaration that
   /// was found by unwrapping an optional context type.
   DeclViaUnwrappedOptional,
-  /// The overload choice materializes a pack from a tuple using
-  /// the \c .element syntax.
+  /// The overload choice materializes a pack from a tuple.
   MaterializePack,
   /// The overload choice indexes into a tuple. Index zero will
   /// have the value of this enumerator, index one will have the value of this

lib/AST/ASTPrinter.cpp

@@ -5223,7 +5223,6 @@ void PrintAST::visitPackExpansionExpr(PackExpansionExpr *expr) {
 
 void PrintAST::visitMaterializePackExpr(MaterializePackExpr *expr) {
   visit(expr->getFromExpr());
-  Printer << ".element";
 }
 
 void PrintAST::visitPackElementExpr(PackElementExpr *expr) {

lib/Sema/CSApply.cpp

@@ -3846,6 +3846,21 @@ namespace {
     }
 
     Expr *visitPackElementExpr(PackElementExpr *expr) {
+      if (auto *packRefExpr = expr->getPackRefExpr()) {
+        packRefExpr = cs.coerceToRValue(packRefExpr);
+        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 *materializedPackExpr = MaterializePackExpr::create(
+              cs.getASTContext(), packRefExpr, packRefExpr->getLoc(),
+              patternType, /*implicit*/ true);
+          cs.cacheType(materializedPackExpr);
+          expr->setPackRefExpr(materializedPackExpr);
+        }
+      }
       return simplifyExprType(expr);
     }
 

lib/Sema/CSGen.cpp

@@ -3148,10 +3148,11 @@ namespace {
             if (type->is<ElementArchetypeType>() &&
                 CS.hasFixFor(CS.getConstraintLocator(declRef),
                              FixKind::IgnoreMissingEachKeyword)) {
-              Packs.push_back(PackElementExpr::create(CS.getASTContext(),
-                                                      /*eachLoc=*/SourceLoc(),
-                                                      declRef,
-                                                      /*implicit=*/true));
+              auto *packElementExpr =
+                  PackElementExpr::create(CS.getASTContext(),
+                                          /*eachLoc=*/{}, declRef,
+                                          /*implicit=*/true, type);
+              Packs.push_back(CS.cacheType(packElementExpr));
             }
           }
 
@@ -3230,8 +3231,17 @@ namespace {
     }
 
     Type visitPackElementExpr(PackElementExpr *expr) {
-      return openPackElement(CS.getType(expr->getPackRefExpr()),
-                             CS.getConstraintLocator(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);
+      }
+
+      return openPackElement(packType, CS.getConstraintLocator(expr));
     }
 
     Type visitMaterializePackExpr(MaterializePackExpr *expr) {

lib/Sema/CSSimplify.cpp

@@ -127,6 +127,14 @@ static bool isPackExpansionType(Type type) {
   return false;
 }
 
+bool constraints::isSingleUnlabeledPackExpansionTuple(Type type) {
+  auto *tuple = type->getRValueType()->getAs<TupleType>();
+  // TODO: drop no name requirement
+  return tuple && (tuple->getNumElements() == 1) &&
+         isPackExpansionType(tuple->getElementType(0)) &&
+         !tuple->getElement(0).hasName();
+}
+
 static bool containsPackExpansionType(ArrayRef<AnyFunctionType::Param> params) {
   return llvm::any_of(params, [&](const auto &param) {
     return isPackExpansionType(param.getPlainType());
@@ -9121,6 +9129,16 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
       return SolutionKind::Solved;
   }
 
+  if (isSingleUnlabeledPackExpansionTuple(patternType)) {
+    auto *elementVar =
+        createTypeVariable(getConstraintLocator(locator), /*options=*/0);
+    addValueMemberConstraint(
+        patternType, DeclNameRef(getASTContext().Id_element), elementVar, DC,
+        FunctionRefKind::Unapplied, {},
+        getConstraintLocator(locator, {ConstraintLocator::Member}));
+    patternType = elementVar;
+  }
+
   // Let's try to resolve element type based on the pattern type.
   if (!patternType->hasTypeVariable()) {
     auto *loc = getConstraintLocator(locator);
@@ -9338,6 +9356,7 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
     if (!memberName.isSpecial()) {
       StringRef nameStr = memberName.getBaseIdentifier().str();
       // Accessing `.element` on an abstract tuple materializes a pack.
+      // (deprecated behavior)
       if (nameStr == "element" && baseTuple->getNumElements() == 1 &&
           isPackExpansionType(baseTuple->getElementType(0))) {
         auto elementType = baseTuple->getElementType(0);
@@ -13312,6 +13331,12 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
     return SolutionKind::Solved;
   }
 
+  // Map element archetypes to the pack context to check for equality.
+  if (packTy->hasElementArchetype()) {
+    auto *packEnv = DC->getGenericEnvironmentOfContext();
+    packTy = packEnv->mapElementTypeIntoPackContext(packTy);
+  }
+
   auto shape = packTy->getReducedShape();
   addConstraint(ConstraintKind::Bind, shapeTy, shape, locator);
   return SolutionKind::Solved;

lib/Sema/ConstraintSystem.cpp

@@ -3559,7 +3559,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     break;
 
   case OverloadChoiceKind::MaterializePack: {
-    // Since `.element` is only applicable to single element tuples at the
+    // Since pack expansion is only applicable to single element tuples at the
     // moment we can just look through l-value base to load it.
     //
     // In the future, _if_ the syntax allows for multiple expansions

test/Constraints/pack-expansion-expressions.swift

@@ -129,14 +129,18 @@ func expansionOfNonPackType<T>(_ t: repeat each T) {}
 
 func tupleExpansion<each T, each U>(
   _ tuple1: (repeat each T),
-  _ tuple2: (repeat each U)
+  _ tuple2: (repeat each U),
+  _ tuple3: inout (repeat each T)
 ) {
-  _ = forward(repeat each tuple1.element)
+  _ = forward(repeat each tuple1)
 
-  _ = zip(repeat each tuple1.element, with: repeat each tuple1.element)
+  _ = zip(repeat each tuple1, with: repeat each tuple1)
+  _ = zip(repeat each tuple1, with: repeat each tuple1.element) // legacy syntax
 
-  _ = zip(repeat each tuple1.element, with: repeat each tuple2.element)
+  _ = zip(repeat each tuple1, with: repeat each tuple2)
   // expected-error@-1 {{global function 'zip(_:with:)' requires the type packs 'each T' and 'each U' have the same shape}}
+
+  _ = forward(repeat each tuple3)
 }
 
 protocol Generatable {
@@ -243,10 +247,10 @@ func test_pack_expansion_materialization_from_lvalue_base() {
 
     init() {
       self.data = (repeat Data<each T>())
-      _ = (repeat each data.element) // Ok
+      _ = (repeat each data) // Ok
 
       var tmp = (repeat Data<each T>()) // expected-warning {{never mutated}}
-      _ = (repeat each tmp.element) // Ok
+      _ = (repeat each tmp) // Ok
 
       // TODO: Add subscript test-case when syntax is supported.
     }
@@ -274,10 +278,9 @@ func packOutsideExpansion<each T>(_ t: repeat each T) {
 
   let tuple = (repeat each t)
 
-  _ = tuple.element
-  // expected-error@-1{{pack reference 'each T' can only appear in pack expansion}}
+  _ = tuple
 
-  _ = each tuple.element
+  _ = each tuple
   // expected-error@-1{{pack reference 'each T' can only appear in pack expansion}}
 }
 
@@ -481,7 +484,6 @@ do {
   func test_misplaced_each<each T: P>(_ value: repeat each T) -> (repeat each T.A) {
     return (repeat each value.makeA())
     // expected-error@-1 {{value pack 'each T' must be referenced with 'each'}} {{25-25=(each }} {{30-30=)}}
-    // expected-error@-2 {{pack expansion requires that '()' and 'each T' have the same shape}}
   }
 }
 
@@ -539,3 +541,15 @@ do {
     test2(repeat each t2, repeat each t1) // expected-error {{local function 'test2' requires that 'each T2' conform to 'RawRepresentable'}}
   }
 }
+
+do {
+  func overloaded<each T>(_ a: Int, _ b: repeat each T) -> (repeat each T) {
+    return (repeat each b)
+  }
+
+  func overloaded() {}
+
+  func test<each T>(_ a: repeat each T) {
+    _ = (repeat overloaded(1, each a))
+  }
+}