[CSClosure] Fix per-element variable finder to correctly handle retur…

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

@@ -365,3 +365,28 @@ func test_diagnosing_on_missing_member_in_case() {
     }
   }
 }
+
+// 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
+  }
+}