Merge pull request #67594 from simanerush/simanerush/pack-iteration-impl

[SE-0408] Enable Pack Iteration

Author: simanerush

include/swift/AST/Decl.h

@@ -5863,6 +5863,7 @@ enum class PropertyWrapperSynthesizedPropertyKind {
 class VarDecl : public AbstractStorageDecl {
   friend class NamingPatternRequest;
   NamedPattern *NamingPattern = nullptr;
+  GenericEnvironment *OpenedElementEnvironment = nullptr;
 
 public:
   enum class Introducer : uint8_t {
@@ -5990,6 +5991,13 @@ class VarDecl : public AbstractStorageDecl {
   NamedPattern *getNamingPattern() const;
   void setNamingPattern(NamedPattern *Pat);
 
+  GenericEnvironment *getOpenedElementEnvironment() const {
+    return OpenedElementEnvironment;
+  }
+  void setOpenedElementEnvironment(GenericEnvironment *Env) {
+    OpenedElementEnvironment = Env;
+  }
+
   /// If this is a VarDecl that does not belong to a CaseLabelItem's pattern,
   /// return this. Otherwise, this VarDecl must belong to a CaseStmt's
   /// CaseLabelItem. In that case, return the first case label item of the first

include/swift/AST/DiagnosticsSema.def

@@ -5859,8 +5859,8 @@ ERROR(expansion_not_allowed,none,
       "pack expansion %0 can only appear in a function parameter list, "
       "tuple element, or generic argument of a variadic type", (Type))
 ERROR(expansion_expr_not_allowed,none,
-      "value pack expansion can only appear inside a function argument list "
-      "or tuple element", ())
+      "value pack expansion can only appear inside a function argument list, "
+      "tuple element, or as the expression of a for-in loop", ())
 ERROR(invalid_expansion_argument,none,
       "cannot pass value pack expansion to non-pack parameter of type %0",
       (Type))
@@ -7730,5 +7730,12 @@ ERROR(referencebindings_binding_must_be_to_lvalue,none,
 ERROR(result_depends_on_no_result,none,
       "Incorrect use of %0 with no result", (StringRef))
 
+//------------------------------------------------------------------------------
+// MARK: Pack Iteration Diagnostics
+//------------------------------------------------------------------------------
+
+ERROR(pack_iteration_where_clause_not_supported, none,
+      "'where' clause in pack iteration is not supported", ())
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/Sema/CSFix.h

@@ -440,6 +440,9 @@ enum class FixKind : uint8_t {
   /// Allow pack expansion expressions in a context that does not support them.
   AllowInvalidPackExpansion,
 
+  /// Ignore `where` clause in a for-in loop with a pack expansion expression.
+  IgnoreWhereClauseInPackIteration,
+
   /// Allow a pack expansion parameter of N elements to be matched
   /// with a single tuple literal argument of the same arity.
   DestructureTupleToMatchPackExpansionParameter,
@@ -2227,6 +2230,26 @@ class AllowInvalidPackExpansion final : public ConstraintFix {
   }
 };
 
+class IgnoreWhereClauseInPackIteration final : public ConstraintFix {
+  IgnoreWhereClauseInPackIteration(ConstraintSystem &cs,
+                                   ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreWhereClauseInPackIteration, locator) {}
+
+public:
+  std::string getName() const override {
+    return "ignore where clause in pack iteration";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  static IgnoreWhereClauseInPackIteration *create(ConstraintSystem &cs,
+                                                  ConstraintLocator *locator);
+
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::IgnoreWhereClauseInPackIteration;
+  }
+};
+
 class CollectionElementContextualMismatch final
     : public ContextualMismatch,
       private llvm::TrailingObjects<CollectionElementContextualMismatch,

include/swift/Sema/ConstraintSystem.h

@@ -3765,6 +3765,11 @@ class ConstraintSystem {
                                        RememberChoice_t rememberChoice,
                                        ConstraintLocatorBuilder locator,
                                        ConstraintFix *compatFix = nullptr);
+  
+  /// Add a materialize constraint for a pack expansion.
+  TypeVariableType *
+  addMaterializePackExpansionConstraint(Type patternType,
+                                        ConstraintLocatorBuilder locator);
 
   /// Add a disjunction constraint.
   void

include/swift/Sema/SyntacticElementTarget.h

@@ -22,15 +22,16 @@
 #include "swift/AST/Pattern.h"
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeLoc.h"
+#include "swift/Basic/TaggedUnion.h"
 #include "swift/Sema/ConstraintLocator.h"
 #include "swift/Sema/ContextualTypeInfo.h"
 
 namespace swift {
 
 namespace constraints {
-/// Describes information about a for-each loop that needs to be tracked
-/// within the constraint system.
-struct ForEachStmtInfo {
+/// Describes information specific to a sequence
+/// in a for-each loop.
+struct SequenceIterationInfo {
   /// The type of the sequence.
   Type sequenceType;
 
@@ -47,6 +48,19 @@ struct ForEachStmtInfo {
   Expr *nextCall;
 };
 
+/// Describes information specific to a pack expansion expression
+/// in a for-each loop.
+struct PackIterationInfo {
+  /// The type of the pattern that matches the elements.
+  Type patternType;
+};
+
+/// Describes information about a for-each loop that needs to be tracked
+/// within the constraint system.
+struct ForEachStmtInfo : TaggedUnion<SequenceIterationInfo, PackIterationInfo> {
+  using TaggedUnion::TaggedUnion;
+};
+
 /// Describes the target to which a constraint system's solution can be
 /// applied.
 class SyntacticElementTarget {

lib/AST/ASTVerifier.cpp

@@ -803,6 +803,13 @@ class Verifier : public ASTWalker {
       if (!shouldVerify(cast<Stmt>(S)))
         return false;
 
+      if (auto *expansion =
+              dyn_cast<PackExpansionExpr>(S->getParsedSequence())) {
+        if (!shouldVerify(expansion)) {
+          return false;
+        }
+      }
+
       if (!S->getElementExpr())
         return true;
 
@@ -812,6 +819,11 @@ class Verifier : public ASTWalker {
     }
 
     void cleanup(ForEachStmt *S) {
+      if (auto *expansion =
+              dyn_cast<PackExpansionExpr>(S->getParsedSequence())) {
+        cleanup(expansion);
+      }
+
       if (!S->getElementExpr())
         return;
 
@@ -2632,6 +2644,16 @@ class Verifier : public ASTWalker {
         abort();
       }
 
+      // If we are performing pack iteration, variables have to carry the
+      // generic environment. Catching the missing environment here will prevent
+      // the code from being lowered.
+      if (var->getTypeInContext()->is<ErrorType>()) {
+        Out << "VarDecl is missing a Generic Environment: ";
+        var->getInterfaceType().print(Out);
+        Out << "\n";
+        abort();
+      }
+
       // The fact that this is *directly* be a reference storage type
       // cuts the code down quite a bit in getTypeOfReference.
       if (var->getAttrs().hasAttribute<ReferenceOwnershipAttr>() !=

lib/AST/ASTWalker.cpp

@@ -1905,11 +1905,9 @@ Stmt *Traversal::visitForEachStmt(ForEachStmt *S) {
   //
   // If for-in is already type-checked, the type-checked version
   // of the sequence is going to be visited as part of `iteratorVar`.
-  if (S->getTypeCheckedSequence()) {
-    if (auto IteratorVar = S->getIteratorVar()) {
-      if (doIt(IteratorVar))
-        return nullptr;
-    }
+  if (auto IteratorVar = S->getIteratorVar()) {
+    if (doIt(IteratorVar))
+      return nullptr;
 
     if (auto NextCall = S->getNextCall()) {
       if ((NextCall = doIt(NextCall)))

lib/AST/Decl.cpp

@@ -7149,6 +7149,11 @@ VarDecl::VarDecl(DeclKind kind, bool isStatic, VarDecl::Introducer introducer,
 }
 
 Type VarDecl::getTypeInContext() const {
+  // If we are performing pack iteration, use the generic environment of the
+  // pack expansion expression to get the right context of a local variable.
+  if (auto *env = getOpenedElementEnvironment())
+    return GenericEnvironment::mapTypeIntoContext(env, getInterfaceType());
+
   return getDeclContext()->mapTypeIntoContext(getInterfaceType());
 }
 

lib/AST/GenericSignature.cpp

@@ -105,8 +105,16 @@ void GenericSignatureImpl::forEachParam(
 
   for (auto req : getRequirements()) {
     GenericTypeParamType *gp;
+    bool isCanonical = false;
     switch (req.getKind()) {
     case RequirementKind::SameType: {
+      if (req.getSecondType()->isParameterPack() != 
+          req.getFirstType()->isParameterPack()) {
+        // This is a same-element requirement, which does not make
+        // type parameters non-canonical.
+        isCanonical = true;
+      }
+
       if (auto secondGP = req.getSecondType()->getAs<GenericTypeParamType>()) {
         // If two generic parameters are same-typed, then the right-hand one
         // is non-canonical.
@@ -136,7 +144,7 @@ void GenericSignatureImpl::forEachParam(
     }
 
     unsigned index = GenericParamKey(gp).findIndexIn(genericParams);
-    genericParamsAreCanonical[index] = false;
+    genericParamsAreCanonical[index] = isCanonical;
   }
 
   // Call the callback with each parameter and the result of the above analysis.

lib/AST/Stmt.cpp

@@ -446,6 +446,9 @@ void ForEachStmt::setPattern(Pattern *p) {
 }
 
 Expr *ForEachStmt::getTypeCheckedSequence() const {
+  if (auto *expansion = dyn_cast<PackExpansionExpr>(getParsedSequence()))
+    return expansion;
+
   return iteratorVar ? iteratorVar->getInit(/*index=*/0) : nullptr;
 }