[ConstraintSystem] Detect and diagnose missing 'each' and provide a fix-it

Diagnose situations where value pack is referenced without an explicit 'each':

```swift
func compute<each T>(_: repeat each T) {}

func test<each T>(v: repeat each T) {
  repeat compute(v) // should be `repeat compute(each v)`
}
```

(cherry picked from commit 31e0a52)
(cherry picked from commit 5836490)
(cherry picked from commit eee7e18)
(cherry picked from commit e3d1317)

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -5435,6 +5435,9 @@ ERROR(pack_reference_must_be_in_expansion,none,
 ERROR(pack_type_requires_keyword_each,none,
       "type pack %0 must be referenced with 'each'",
       (TypeRepr*))
+ERROR(value_pack_requires_keyword_each,none,
+      "value pack %0 must be referenced with 'each'",
+      (Type))
 ERROR(tuple_duplicate_label,none,
       "cannot create a tuple with a duplicate element label", ())
 ERROR(multiple_ellipsis_in_tuple,none,

include/swift/AST/Expr.h

@@ -3627,8 +3627,6 @@ class PackExpansionExpr final : public Expr {
     PatternExpr = patternExpr;
   }
 
-  void getExpandedPacks(SmallVectorImpl<ASTNode> &packs);
-
   GenericEnvironment *getGenericEnvironment() {
     return Environment;
   }

include/swift/Sema/CSFix.h

@@ -441,6 +441,9 @@ enum class FixKind : uint8_t {
 
   /// Allow value pack expansion without pack references.
   AllowValueExpansionWithoutPackReferences,
+
+  /// Ignore missing 'each' keyword before value pack reference.
+  IgnoreMissingEachKeyword,
 };
 
 class ConstraintFix {
@@ -3481,6 +3484,33 @@ class AllowValueExpansionWithoutPackReferences final : public ConstraintFix {
   }
 };
 
+class IgnoreMissingEachKeyword final : public ConstraintFix {
+  Type ValuePackType;
+
+  IgnoreMissingEachKeyword(ConstraintSystem &cs, Type valuePackTy,
+                           ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreMissingEachKeyword, locator),
+        ValuePackType(valuePackTy) {}
+
+public:
+  std::string getName() const override {
+    return "allow value pack reference without 'each'";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  bool diagnoseForAmbiguity(CommonFixesArray commonFixes) const override {
+    return diagnose(*commonFixes.front().first);
+  }
+
+  static IgnoreMissingEachKeyword *
+  create(ConstraintSystem &cs, Type valuePackTy, ConstraintLocator *locator);
+
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::IgnoreMissingEachKeyword;
+  }
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/AST/Expr.cpp

@@ -1279,46 +1279,6 @@ PackExpansionExpr::create(ASTContext &ctx, SourceLoc repeatLoc,
                                      implicit, type);
 }
 
-void PackExpansionExpr::getExpandedPacks(SmallVectorImpl<ASTNode> &packs) {
-  struct PackCollector : public ASTWalker {
-    llvm::SmallVector<ASTNode, 2> packs;
-
-    /// Walk everything that's available.
-    MacroWalking getMacroWalkingBehavior() const override {
-      return MacroWalking::ArgumentsAndExpansion;
-    }
-
-    virtual PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-      // Don't walk into nested pack expansions
-      if (isa<PackExpansionExpr>(E)) {
-        return Action::SkipChildren(E);
-      }
-
-      if (isa<PackElementExpr>(E)) {
-        packs.push_back(E);
-      }
-
-      return Action::Continue(E);
-    }
-
-    virtual PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
-      // Don't walk into nested pack expansions
-      if (isa<PackExpansionTypeRepr>(T)) {
-        return Action::SkipChildren();
-      }
-
-      if (isa<PackElementTypeRepr>(T)) {
-        packs.push_back(T);
-      }
-
-      return Action::Continue();
-    }
-  } packCollector;
-
-  getPatternExpr()->walk(packCollector);
-  packs.append(packCollector.packs.begin(), packCollector.packs.end());
-}
-
 PackElementExpr *
 PackElementExpr::create(ASTContext &ctx, SourceLoc eachLoc, Expr *packRefExpr,
                         bool implicit, Type type) {

lib/Sema/CSDiagnostics.cpp

@@ -8904,3 +8904,25 @@ bool ValuePackExpansionWithoutPackReferences::diagnoseAsError() {
   emitDiagnostic(diag::value_expansion_not_variadic);
   return true;
 }
+
+bool MissingEachForValuePackReference::diagnoseAsError() {
+  bool fixItNeedsParens = false;
+  // If 'each' is missing form a base of a member reference
+  // it has to be wrapped in parens.
+  if (auto anchor = getAsExpr(getAnchor())) {
+    fixItNeedsParens = isExpr<UnresolvedDotExpr>(findParentExpr(anchor));
+  }
+
+  {
+    auto diagnostic = emitDiagnostic(diag::value_pack_requires_keyword_each, ValuePackType);
+    if (fixItNeedsParens) {
+      auto range = getSourceRange();
+      diagnostic.fixItInsert(range.Start, "(each ")
+          .fixItInsertAfter(range.End, ")");
+    } else {
+      diagnostic.fixItInsert(getLoc(), "each ");
+    }
+  }
+
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -3017,6 +3017,27 @@ class ValuePackExpansionWithoutPackReferences final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+/// Diagnose situations where value pack is referenced without explicit 'each':
+///
+/// \code
+/// func compute<each T>(_: repeat each T) {}
+///
+/// func test<each T>(v: repeat each T) {
+///   repeat compute(v) // should be `repeat compute(each v)`
+/// }
+/// \endcode
+class MissingEachForValuePackReference final : public FailureDiagnostic {
+  Type ValuePackType;
+
+public:
+  MissingEachForValuePackReference(const Solution &solution, Type valuePackTy,
+                                   ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator),
+        ValuePackType(resolveType(valuePackTy)) {}
+
+  bool diagnoseAsError() override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -2769,3 +2769,17 @@ AllowValueExpansionWithoutPackReferences::create(ConstraintSystem &cs,
   return new (cs.getAllocator())
       AllowValueExpansionWithoutPackReferences(cs, locator);
 }
+
+bool IgnoreMissingEachKeyword::diagnose(const Solution &solution,
+                                        bool asNote) const {
+  MissingEachForValuePackReference failure(solution, ValuePackType,
+                                           getLocator());
+  return failure.diagnose(asNote);
+}
+
+IgnoreMissingEachKeyword *
+IgnoreMissingEachKeyword::create(ConstraintSystem &cs, Type valuePackTy,
+                                 ConstraintLocator *locator) {
+  return new (cs.getAllocator())
+      IgnoreMissingEachKeyword(cs, valuePackTy, locator);
+}

lib/Sema/CSGen.cpp

@@ -1089,6 +1089,23 @@ 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()) {
+        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());
+      return openPackElement(packType, /*packRepr*/ nullptr);
+    }
+
   public:
     ConstraintGenerator(ConstraintSystem &CS, DeclContext *DC)
         : CS(CS), CurDC(DC ? DC : CS.DC), CurrPhase(CS.getPhase()) {
@@ -1407,6 +1424,20 @@ namespace {
             return invalidateReference();
           }
 
+          // value packs cannot be referenced without `each` immediately
+          // preceding them.
+          if (auto *expansion = knownType->getAs<PackExpansionType>()) {
+            if (!PackElementEnvironments.empty() &&
+                !isExpr<PackElementExpr>(CS.getParentExpr(E))) {
+              auto packType = expansion->getPatternType();
+              (void)CS.recordFix(
+                  IgnoreMissingEachKeyword::create(CS, packType, locator));
+              auto eltType = openPackElement(packType, locator);
+              CS.setType(E, eltType);
+              return eltType;
+            }
+          }
+
           if (!knownType->hasPlaceholder()) {
             // Set the favored type for this expression to the known type.
             CS.setFavoredType(E, knownType.getPointer());
@@ -3057,6 +3088,67 @@ namespace {
       return variadicSeq;
     }
 
+    void collectExpandedPacks(PackExpansionExpr *expansion,
+                              SmallVectorImpl<ASTNode> &packs) {
+      struct PackCollector : public ASTWalker {
+      private:
+        ConstraintSystem &CS;
+        SmallVectorImpl<ASTNode> &Packs;
+
+      public:
+        PackCollector(ConstraintSystem &cs, SmallVectorImpl<ASTNode> &packs)
+            : CS(cs), Packs(packs) {}
+
+        /// Walk everything that's available.
+        MacroWalking getMacroWalkingBehavior() const override {
+          return MacroWalking::ArgumentsAndExpansion;
+        }
+
+        virtual PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+          // Don't walk into nested pack expansions
+          if (isa<PackExpansionExpr>(E)) {
+            return Action::SkipChildren(E);
+          }
+
+          if (isa<PackElementExpr>(E)) {
+            Packs.push_back(E);
+          }
+
+          if (auto *declRef = dyn_cast<DeclRefExpr>(E)) {
+            auto type = CS.getTypeIfAvailable(declRef);
+            if (!type)
+              return Action::Continue(E);
+
+            if (type->is<ElementArchetypeType>() &&
+                CS.hasFixFor(CS.getConstraintLocator(declRef),
+                             FixKind::IgnoreMissingEachKeyword)) {
+              Packs.push_back(PackElementExpr::create(CS.getASTContext(),
+                                                      /*eachLoc=*/SourceLoc(),
+                                                      declRef,
+                                                      /*implicit=*/true));
+            }
+          }
+
+          return Action::Continue(E);
+        }
+
+        virtual PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
+          // Don't walk into nested pack expansions
+          if (isa<PackExpansionTypeRepr>(T)) {
+            return Action::SkipChildren();
+          }
+
+          if (isa<PackElementTypeRepr>(T)) {
+            Packs.push_back(T);
+          }
+
+          return Action::Continue();
+        }
+      } packCollector(CS, packs);
+
+      expansion->getPatternExpr()->walk(packCollector);
+    }
+
     Type visitPackExpansionExpr(PackExpansionExpr *expr) {
       assert(PackElementEnvironments.back() == expr);
       PackElementEnvironments.pop_back();
@@ -3080,7 +3172,7 @@ namespace {
       // Generate ShapeOf constraints between all packs expanded by this
       // pack expansion expression through the shape type variable.
       SmallVector<ASTNode, 2> expandedPacks;
-      expr->getExpandedPacks(expandedPacks);
+      collectExpandedPacks(expr, expandedPacks);
 
       if (expandedPacks.empty()) {
         (void)CS.recordFix(AllowValueExpansionWithoutPackReferences::create(
@@ -3112,23 +3204,8 @@ namespace {
     }
 
     Type visitPackElementExpr(PackElementExpr *expr) {
-      auto packType = CS.getType(expr->getPackRefExpr());
-
-      // 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()) {
-        return CS.createTypeVariable(CS.getConstraintLocator(expr),
-                                     TVO_CanBindToHole |
-                                     TVO_CanBindToNoEscape);
-      }
-
-      // The type of a PackElementExpr is the opened pack element archetype
-      // of the pack reference.
-      OpenPackElementType openPackElement(CS, CS.getConstraintLocator(expr),
-                                          PackElementEnvironments.back());
-      return openPackElement(packType, /*packRepr*/ nullptr);
+      return openPackElement(CS.getType(expr->getPackRefExpr()),
+                             CS.getConstraintLocator(expr));
     }
 
     Type visitMaterializePackExpr(MaterializePackExpr *expr) {

lib/Sema/CSSimplify.cpp

@@ -14822,6 +14822,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AddExplicitExistentialCoercion:
   case FixKind::DestructureTupleToMatchPackExpansionParameter:
   case FixKind::AllowValueExpansionWithoutPackReferences:
+  case FixKind::IgnoreMissingEachKeyword:
     llvm_unreachable("handled elsewhere");
   }
 

test/Constraints/pack-expansion-expressions.swift

@@ -454,7 +454,8 @@ do {
 
   func test_misplaced_each<each T: P>(_ value: repeat each T) -> (repeat each T.A) {
     return (repeat each value.makeA())
-    // expected-error@-1 {{pack reference 'each T' can only appear in pack expansion}}
+    // 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}}
   }
 }
 
@@ -480,3 +481,16 @@ do {
     // expected-error@-1:25 {{value pack expansion must contain at least one pack reference}}
   }
 }
+
+// missing 'each' keyword before value pack references
+do {
+  func overloaded<each U>(_: String, _: repeat each U) -> Int { 42 }
+  func overloaded<each T>(_: Int, _ b: repeat each T) -> (repeat each T) {
+    fatalError()
+  }
+
+  func test<each T>(v: repeat each T) {
+    _ = (repeat overloaded(42, v)) // expected-error {{value pack 'each T' must be referenced with 'each'}} {{32-32=each }}
+    _ = (repeat overloaded(42, each v)) // Ok
+  }
+}