[Constraint generation] Collect all closure body references to type variables.

Collect all references to parameters whose types involve type variables,
including in closures that aren’t single-expression. This fixes a type
checker assertion that occurs when the constraint graph can get disconnected
with the combination of delayed constraint generation for single-expression
closures and the use of function builders.

Fixes rdar://problem/58695803.

Author: DougGregor

lib/Sema/CSGen.cpp

@@ -2455,43 +2455,43 @@ namespace {
       return tryFinder.foundThrow();
     }
 
-    void collectParameterRefs(Expr *expr,
-                              llvm::SmallVectorImpl<TypeVariableType *> &refs) {
-      expr->forEachChildExpr([&](Expr *childExpr) -> Expr * {
-        if (auto *closure = dyn_cast<ClosureExpr>(childExpr)) {
-          if (closure->hasSingleExpressionBody()) {
-            collectParameterRefs(closure->getSingleExpressionBody(), refs);
-            return childExpr;
-          }
-        }
-
-        if (auto *DRE = dyn_cast<DeclRefExpr>(childExpr)) {
-          if (auto *PD = dyn_cast<ParamDecl>(DRE->getDecl())) {
-            if (CS.hasType(PD)) {
-              if (auto *paramType = CS.getType(PD)->getAs<TypeVariableType>())
-                refs.push_back(paramType);
-            }
-          }
-        }
-
-        return childExpr;
-      });
-    }
-
     Type visitClosureExpr(ClosureExpr *closure) {
       auto *locator = CS.getConstraintLocator(closure);
       auto closureType = CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
 
-      llvm::SmallVector<TypeVariableType *, 4> paramRefs;
-      collectParameterRefs(closure, paramRefs);
+      // Collect any references to closure parameters whose types involve type
+      // variables from the closure, because there will be a dependency on
+      // those type variables once we have generated constraints for the
+      // closure body.
+      struct CollectParameterRefs : public ASTWalker {
+        ConstraintSystem &cs;
+        llvm::SmallVector<TypeVariableType *, 4> paramRefs;
+
+        CollectParameterRefs(ConstraintSystem &cs) : cs(cs) { }
+
+        std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
+          // Retrieve type variables from references to parameter declarations.
+          if (auto *declRef = dyn_cast<DeclRefExpr>(expr)) {
+            if (auto *paramDecl = dyn_cast<ParamDecl>(declRef->getDecl())) {
+              if (Type paramType = cs.getTypeIfAvailable(paramDecl)) {
+                paramType->getTypeVariables(paramRefs);
+              }
+            }
+          }
+
+          return { true, expr };
+        }
+      } collectParameterRefs(CS);
+      closure->walk(collectParameterRefs);
 
       auto inferredType = inferClosureType(closure);
       if (!inferredType || inferredType->hasError())
         return Type();
 
       CS.addUnsolvedConstraint(
           Constraint::create(CS, ConstraintKind::DefaultClosureType,
-                             closureType, inferredType, locator, paramRefs));
+                             closureType, inferredType, locator,
+                             collectParameterRefs.paramRefs));
 
       CS.setClosureType(closure, inferredType);
       return closureType;

test/Constraints/function_builder.swift

@@ -8,6 +8,10 @@ enum Either<T,U> {
 
 @_functionBuilder
 struct TupleBuilder {
+  static func buildBlock<T1>(_ t1: T1) -> (T1) {
+    return (t1)
+  }
+
   static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
     return (t1, t2)
   }
@@ -430,3 +434,30 @@ func test_single_stmt_closure_support() {
 
   let _ = test { 0 } // ok
 }
+
+// Check a case involving nested closures that refer to parameters of their
+// enclosing closures.
+struct X<C: Collection, T> {
+  init(_ c: C, @TupleBuilder body: (C.Element) -> T) { }
+}
+
+struct Y<T> {
+  init(@TupleBuilder body: () -> T) { }
+}
+
+struct Z<T> {
+  init(@TupleBuilder body: () -> T) { }
+}
+
+func testNestedClosuresWithDependencies(cond: Bool) {
+  tuplify(cond) { _ in
+    X([1, 2, 3]) { x in
+      Y {
+        Z {
+          x
+          1
+        }
+      }
+    }
+  }
+}