Merge pull request #73849 from xedin/collect-closures-to-pack-expansion-they-appear-in

[ConstraintSystem] Fix type-checking of closures in pack expansion contexts

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);
 }
 

lib/Sema/CSSyntacticElement.cpp

@@ -153,6 +153,17 @@ class TypeVariableRefFinder : public ASTWalker {
       }
     }
 
+    // If closure appears inside of a pack expansion, the elements
+    // that reference pack elements have to bring expansion's shape
+    // type in scope to make sure that the shapes match.
+    if (auto *packElement = getAsExpr<PackElementExpr>(expr)) {
+      if (auto *outerEnvironment = CS.getPackEnvironment(packElement)) {
+        auto *expansionTy = CS.simplifyType(CS.getType(outerEnvironment))
+                                ->castTo<PackExpansionType>();
+        expansionTy->getCountType()->getTypeVariables(ReferencedVars);
+      }
+    }
+
     return Action::Continue(expr);
   }
 

test/Constraints/pack-expansion-expressions.swift

@@ -739,3 +739,25 @@ 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
+  }
+
+  func testMultiStmtClosure<each T>(_ t: repeat each T) -> (repeat each T) {
+     (repeat takesClosure {
+       let v = each t
+       return v
+    }) // Ok
+  }
+
+  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
+  }
+}