Merge pull request #41639 from xedin/se-0326-special-handling-of-local-funcs

[CSClosure] Delay type-checking of local functions until body is rewr…

Author: xedin

lib/Sema/CSClosure.cpp

@@ -1004,6 +1004,9 @@ class ClosureConstraintApplication
   RewriteTargetFn rewriteTarget;
   bool isSingleExpression;
 
+  /// All `func`s declared in the body of the closure.
+  SmallVector<FuncDecl *, 4> LocalFuncs;
+
 public:
   /// Whether an error was encountered while generating constraints.
   bool hadError = false;
@@ -1048,6 +1051,14 @@ class ClosureConstraintApplication
       // information e.g. accessors and do access/availability checks.
     }
 
+    // Local functions cannot be type-checked in-order because they can
+    // capture variables declared after them. Let's save them to be
+    // processed after the solution has been applied to the body.
+    if (auto *func = dyn_cast<FuncDecl>(decl)) {
+      LocalFuncs.push_back(func);
+      return;
+    }
+
     TypeChecker::typeCheckDecl(decl);
   }
 
@@ -1456,6 +1467,19 @@ class ClosureConstraintApplication
   UNSUPPORTED_STMT(Fail)
 #undef UNSUPPORTED_STMT
 
+public:
+  /// Apply solution to the closure and return updated body.
+  ASTNode apply() {
+    auto body = visit(closure->getBody());
+
+    // Since local functions can capture variables that are declared
+    // after them, let's type-check them after all of the pattern
+    // bindings have been resolved by applying solution to the body.
+    for (auto *func : LocalFuncs)
+      TypeChecker::typeCheckDecl(func);
+
+    return body;
+  }
 };
 
 }
@@ -1552,7 +1576,7 @@ bool ConstraintSystem::applySolutionToBody(Solution &solution,
   auto closureType = cs.getType(closure)->castTo<FunctionType>();
   ClosureConstraintApplication application(
       solution, closure, closureType->getResult(), rewriteTarget);
-  auto body = application.visit(closure->getBody());
+  auto body = application.apply();
 
   if (!body || application.hadError)
     return true;

test/expr/closure/multi_statement.swift

@@ -208,3 +208,28 @@ func test_custom_tilde_equals_operator_matching() {
     }
   }
 }
+
+// Local functions can capture variables before they are declared.
+func test_local_function_capturing_vars() {
+  struct A {
+    var cond: Bool
+  }
+
+  func test<T>(fn: () -> T) -> T {
+    fn()
+  }
+
+  func outer(a: A) {
+    test {
+      func local() {
+        if !message.isEmpty { // Ok
+          print(message)
+        }
+
+        message = "World" // Ok
+      }
+
+      var message = a.cond ? "hello" : ""
+    }
+  }
+}