[5.9][Sema] Improvements to variadic generics inference and diagnostics

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

@@ -24,6 +24,7 @@
 #include "swift/AST/DiagnosticsClangImporter.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/Initializer.h"
@@ -88,6 +89,54 @@ Type FailureDiagnostic::getRawType(ASTNode node) const {
   return S.getType(node);
 }
 
+Type FailureDiagnostic::resolveType(Type rawType, bool reconstituteSugar,
+                                    bool wantRValue) const {
+  rawType = rawType.transform([&](Type type) -> Type {
+    if (auto *typeVar = type->getAs<TypeVariableType>()) {
+      auto resolvedType = S.simplifyType(typeVar);
+
+      if (!resolvedType->hasUnresolvedType())
+        return resolvedType;
+
+      // If type variable was simplified to an unresolved pack expansion
+      // type, let's examine its original pattern type because it could
+      // contain type variables replaceable with their generic parameter
+      // types.
+      if (auto *expansion = resolvedType->getAs<PackExpansionType>()) {
+        auto *locator = typeVar->getImpl().getLocator();
+        auto *openedExpansionTy =
+            locator->castLastElementTo<LocatorPathElt::PackExpansionType>()
+                .getOpenedType();
+        auto patternType = resolveType(openedExpansionTy->getPatternType());
+        return PackExpansionType::get(patternType, expansion->getCountType());
+      }
+
+      Type GP = typeVar->getImpl().getGenericParameter();
+      return resolvedType->is<UnresolvedType>() && GP ? GP : resolvedType;
+    }
+
+    if (type->hasElementArchetype()) {
+      auto *env = getDC()->getGenericEnvironmentOfContext();
+      return env->mapElementTypeIntoPackContext(type);
+    }
+
+    if (auto *packType = type->getAs<PackType>()) {
+      if (packType->getNumElements() == 1) {
+        auto eltType = resolveType(packType->getElementType(0));
+        if (auto expansion = eltType->getAs<PackExpansionType>())
+          return expansion->getPatternType();
+      }
+    }
+
+    return type->isPlaceholder() ? Type(type->getASTContext().TheUnresolvedType)
+                                 : type;
+  });
+
+  if (reconstituteSugar)
+    rawType = rawType->reconstituteSugar(/*recursive*/ true);
+  return wantRValue ? rawType->getRValueType() : rawType;
+}
+
 template <typename... ArgTypes>
 InFlightDiagnostic
 FailureDiagnostic::emitDiagnostic(ArgTypes &&... Args) const {
@@ -8904,3 +8953,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

@@ -98,48 +98,7 @@ class FailureDiagnostic {
 
   /// Resolve type variables present in the raw type, if any.
   Type resolveType(Type rawType, bool reconstituteSugar = false,
-                   bool wantRValue = true) const {
-    rawType = rawType.transform([&](Type type) -> Type {
-      if (auto *typeVar = type->getAs<TypeVariableType>()) {
-        auto resolvedType = S.simplifyType(typeVar);
-
-        if (!resolvedType->hasUnresolvedType())
-          return resolvedType;
-
-        // If type variable was simplified to an unresolved pack expansion
-        // type, let's examine its original pattern type because it could
-        // contain type variables replaceable with their generic parameter
-        // types.
-        if (auto *expansion = resolvedType->getAs<PackExpansionType>()) {
-          auto *locator = typeVar->getImpl().getLocator();
-          auto *openedExpansionTy =
-              locator->castLastElementTo<LocatorPathElt::PackExpansionType>()
-                  .getOpenedType();
-          auto patternType = resolveType(openedExpansionTy->getPatternType());
-          return PackExpansionType::get(patternType, expansion->getCountType());
-        }
-
-        Type GP = typeVar->getImpl().getGenericParameter();
-        return resolvedType->is<UnresolvedType>() && GP ? GP : resolvedType;
-      }
-
-      if (auto *packType = type->getAs<PackType>()) {
-        if (packType->getNumElements() == 1) {
-          auto eltType = resolveType(packType->getElementType(0));
-          if (auto expansion = eltType->getAs<PackExpansionType>())
-            return expansion->getPatternType();
-        }
-      }
-
-      return type->isPlaceholder()
-                 ? Type(type->getASTContext().TheUnresolvedType)
-                 : type;
-    });
-
-    if (reconstituteSugar)
-      rawType = rawType->reconstituteSugar(/*recursive*/ true);
-    return wantRValue ? rawType->getRValueType() : rawType;
-  }
+                   bool wantRValue = true) const;
 
   template <typename... ArgTypes>
   InFlightDiagnostic emitDiagnostic(ArgTypes &&... Args) const;
@@ -3017,6 +2976,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);
+}