[CSClosure] Fix per-element variable finder to correctly handle return statements

Type finder is still allowed to walk into closures to find any
referenced variables, but it should bring external result type
into scope only if a particular `return` belongs to the same
closure as the element.

Author: xedin

lib/Sema/CSClosure.cpp

@@ -37,6 +37,9 @@ Expr *getVoidExpr(ASTContext &ctx) {
 
 /// Find any type variable references inside of an AST node.
 class TypeVariableRefFinder : public ASTWalker {
+  /// A stack of all closures the walker encountered so far.
+  SmallVector<DeclContext *> ClosureDCs;
+
   ConstraintSystem &CS;
   ASTNode Parent;
 
@@ -46,9 +49,16 @@ class TypeVariableRefFinder : public ASTWalker {
   TypeVariableRefFinder(
       ConstraintSystem &cs, ASTNode parent,
       llvm::SmallPtrSetImpl<TypeVariableType *> &referencedVars)
-      : CS(cs), Parent(parent), ReferencedVars(referencedVars) {}
+      : CS(cs), Parent(parent), ReferencedVars(referencedVars) {
+    if (auto *closure = getAsExpr<ClosureExpr>(Parent))
+      ClosureDCs.push_back(closure);
+  }
 
   std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
+    if (auto *closure = dyn_cast<ClosureExpr>(expr)) {
+      ClosureDCs.push_back(closure);
+    }
+
     if (auto *DRE = dyn_cast<DeclRefExpr>(expr)) {
       auto *decl = DRE->getDecl();
 
@@ -81,20 +91,33 @@ class TypeVariableRefFinder : public ASTWalker {
     return {true, expr};
   }
 
+  Expr *walkToExprPost(Expr *expr) override {
+    if (auto *closure = dyn_cast<ClosureExpr>(expr)) {
+      ClosureDCs.pop_back();
+    }
+    return expr;
+  }
+
   std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
     // Return statements have to reference outside result type
     // since all of them are joined by it if it's not specified
     // explicitly.
     if (isa<ReturnStmt>(stmt)) {
       if (auto *closure = getAsExpr<ClosureExpr>(Parent)) {
-        inferVariables(CS.getClosureType(closure)->getResult());
+        // Return is only viable if it belongs to a parent closure.
+        if (currentClosureDC() == closure)
+          inferVariables(CS.getClosureType(closure)->getResult());
       }
     }
 
     return {true, stmt};
   }
 
 private:
+  DeclContext *currentClosureDC() const {
+    return ClosureDCs.empty() ? nullptr : ClosureDCs.back();
+  }
+
   void inferVariables(Type type) {
     type = type->getWithoutSpecifierType();
     // Record the type variable itself because it has to

test/expr/closure/multi_statement.swift

@@ -350,3 +350,28 @@ func test_no_crash_with_circular_ref_due_to_error() {
     return 0
   }
 }
+
+// Type finder shouldn't bring external closure result type
+// into the scope of an inner closure e.g. while solving
+// init of pattern binding `x`.
+func test_type_finder_doesnt_walk_into_inner_closures() {
+  func test<T>(fn: () -> T) -> T { fn() }
+
+  _ = test { // Ok
+    let x = test {
+      42
+    }
+
+    let _ = test {
+      test { "" }
+    }
+
+    // multi-statement
+    let _ = test {
+      _ = 42
+      return test { "" }
+    }
+
+    return x
+  }
+}