[ConstraintSystem] Connect closure to outer pack expansions it appears in

Type variable reference collector needs to be augmented to collect
type variables associated with pack expansions that a closure references
elements of, otherwise it would get disconnected from the context.

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -6453,6 +6453,7 @@ class ConjunctionElementProducer : public BindingProducer<ConjunctionElement> {
 ///
 /// This includes:
 /// - Not yet resolved outer VarDecls (including closure parameters)
+/// - Outer pack expansions that are not yet fully resolved
 /// - Return statements with a contextual type that has not yet been resolved
 ///
 /// This is required because isolated conjunctions, just like single-expression
@@ -6474,6 +6475,7 @@ class TypeVarRefCollector : public ASTWalker {
 
   /// Infer the referenced type variables from a given decl.
   void inferTypeVars(Decl *D);
+  void inferTypeVars(PackExpansionExpr *);
 
   MacroWalking getMacroWalkingBehavior() const override {
     return MacroWalking::Arguments;

lib/Sema/CSGen.cpp

@@ -870,6 +870,14 @@ void TypeVarRefCollector::inferTypeVars(Decl *D) {
   TypeVars.insert(typeVars.begin(), typeVars.end());
 }
 
+void TypeVarRefCollector::inferTypeVars(PackExpansionExpr *E) {
+  auto expansionType = CS.getType(E)->castTo<PackExpansionType>();
+
+  SmallPtrSet<TypeVariableType *, 4> referencedVars;
+  expansionType->getTypeVariables(referencedVars);
+  TypeVars.insert(referencedVars.begin(), referencedVars.end());
+}
+
 ASTWalker::PreWalkResult<Expr *>
 TypeVarRefCollector::walkToExprPre(Expr *expr) {
   if (isa<ClosureExpr>(expr))
@@ -891,6 +899,14 @@ TypeVarRefCollector::walkToExprPre(Expr *expr) {
       inferTypeVars(D);
     }
   }
+
+  if (auto *packElement = getAsExpr<PackElementExpr>(expr)) {
+    // If environment hasn't been established yet, it means that pack expansion
+    // appears inside of this closure.
+    if (auto *outerEnvironment = CS.getPackEnvironment(packElement))
+      inferTypeVars(outerEnvironment);
+  }
+
   return Action::Continue(expr);
 }
 

test/Constraints/pack-expansion-expressions.swift

@@ -739,3 +739,27 @@ do {
     // expected-warning@-2 {{immutable value 'y' was never used; consider replacing with '_' or removing it}}
   }
 }
+
+// Closures wrapped in a pack expansion
+do {
+  func takesClosure<T>(_ fn: () -> T) -> T { return fn() }
+
+  func testClosure<each T>(_ t: repeat each T) -> (repeat each T) {
+    (repeat takesClosure { each t }) // Ok
+  }
+
+  // FIXME: multi-statement closures should type-check.
+  func testMultiStmtClosure<each T>(_ t: repeat each T) -> (repeat each T) {
+     (repeat takesClosure {
+       // expected-error@-1 {{pack expansion requires that '_' and 'each T' have the same shape}}
+       let v = each t
+       return v
+    })
+  }
+
+  func takesAutoclosure<T>(_ fn: @autoclosure () -> T) -> T { return fn() }
+
+  func f2<each T>(_ t: repeat each T) -> (repeat each T) {
+    (repeat takesAutoclosure(each t)) // Ok
+  }
+}