[CS] Connect conjunctions with ReturnStmts to return type

Augment the TypeVarRefCollector such that it
picks up any type variables present in the result
type for a closure DeclContext when visiting a
ReturnStmt. This ensures we correctly handle
if/switch expressions that contain `return`
statements.

rdar://114402042

Author: hamishknight

include/swift/Sema/ConstraintSystem.h

@@ -6243,17 +6243,24 @@ class ConjunctionElementProducer : public BindingProducer<ConjunctionElement> {
 ///
 /// This includes:
 /// - Not yet resolved outer VarDecls (including closure parameters)
+/// - Return statements with a contextual type that has not yet been resolved
 ///
 /// This is required because isolated conjunctions, just like single-expression
 /// closures, have to be connected to type variables they are going to use,
 /// otherwise they'll get placed in a separate solver component and would never
 /// produce a solution.
 class TypeVarRefCollector : public ASTWalker {
   ConstraintSystem &CS;
+  DeclContext *DC;
+  ConstraintLocator *Locator;
+
   llvm::SmallSetVector<TypeVariableType *, 4> TypeVars;
+  unsigned DCDepth = 0;
 
 public:
-  TypeVarRefCollector(ConstraintSystem &cs) : CS(cs) {}
+  TypeVarRefCollector(ConstraintSystem &cs, DeclContext *dc,
+                      ConstraintLocator *locator)
+      : CS(cs), DC(dc), Locator(locator) {}
 
   /// Infer the referenced type variables from a given decl.
   void inferTypeVars(Decl *D);
@@ -6263,6 +6270,8 @@ class TypeVarRefCollector : public ASTWalker {
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *expr) override;
+  PostWalkResult<Expr *> walkToExprPost(Expr *expr) override;
+  PreWalkResult<Stmt *> walkToStmtPre(Stmt *stmt) override;
 
   PreWalkAction walkToDeclPre(Decl *D) override {
     // We only need to walk into PatternBindingDecls, other kinds of decls

lib/Sema/CSGen.cpp

@@ -882,6 +882,9 @@ void TypeVarRefCollector::inferTypeVars(Decl *D) {
 
 ASTWalker::PreWalkResult<Expr *>
 TypeVarRefCollector::walkToExprPre(Expr *expr) {
+  if (isa<ClosureExpr>(expr))
+    DCDepth += 1;
+
   if (auto *DRE = dyn_cast<DeclRefExpr>(expr))
     inferTypeVars(DRE->getDecl());
 
@@ -901,6 +904,34 @@ TypeVarRefCollector::walkToExprPre(Expr *expr) {
   return Action::Continue(expr);
 }
 
+ASTWalker::PostWalkResult<Expr *>
+TypeVarRefCollector::walkToExprPost(Expr *expr) {
+  if (isa<ClosureExpr>(expr))
+    DCDepth -= 1;
+
+  return Action::Continue(expr);
+}
+
+ASTWalker::PreWalkResult<Stmt *>
+TypeVarRefCollector::walkToStmtPre(Stmt *stmt) {
+  // If we have a return without any intermediate DeclContexts in a ClosureExpr,
+  // we need to include any type variables in the closure's result type, since
+  // the conjunction will generate constraints using that type. We don't need to
+  // connect to returns in e.g nested closures since we'll connect those when we
+  // generate constraints for those closures. We also don't need to bother if
+  // we're generating constraints for the closure itself, since we'll connect
+  // the conjunction to the closure type variable itself.
+  if (auto *CE = dyn_cast<ClosureExpr>(DC)) {
+    if (isa<ReturnStmt>(stmt) && DCDepth == 0 &&
+        !Locator->directlyAt<ClosureExpr>()) {
+      SmallPtrSet<TypeVariableType *, 4> typeVars;
+      CS.getClosureType(CE)->getResult()->getTypeVariables(typeVars);
+      TypeVars.insert(typeVars.begin(), typeVars.end());
+    }
+  }
+  return Action::Continue(stmt);
+}
+
 namespace {
   class ConstraintGenerator : public ExprVisitor<ConstraintGenerator, Type> {
     ConstraintSystem &CS;
@@ -1304,7 +1335,8 @@ namespace {
           // in the tap body, otherwise tap expression is going
           // to get disconnected from the context.
           if (auto *body = tap->getBody()) {
-            TypeVarRefCollector refCollector(CS);
+            TypeVarRefCollector refCollector(
+                CS, tap->getVar()->getDeclContext(), locator);
 
             body->walk(refCollector);
 
@@ -2938,7 +2970,7 @@ namespace {
       auto *locator = CS.getConstraintLocator(closure);
       auto closureType = CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
 
-      TypeVarRefCollector refCollector(CS);
+      TypeVarRefCollector refCollector(CS, /*DC*/ closure, locator);
       // Walk the capture list if this closure has one,  because it could
       // reference declarations from the outer closure.
       if (auto *captureList =

lib/Sema/CSSyntacticElement.cpp

@@ -281,11 +281,10 @@ static bool isViableElement(ASTNode element,
 using ElementInfo = std::tuple<ASTNode, ContextualTypeInfo,
                                /*isDiscarded=*/bool, ConstraintLocator *>;
 
-static void createConjunction(ConstraintSystem &cs,
+static void createConjunction(ConstraintSystem &cs, DeclContext *dc,
                               ArrayRef<ElementInfo> elements,
-                              ConstraintLocator *locator,
-                              bool isIsolated = false,
-                              ArrayRef<TypeVariableType *> extraTypeVars = {}) {
+                              ConstraintLocator *locator, bool isIsolated,
+                              ArrayRef<TypeVariableType *> extraTypeVars) {
   SmallVector<Constraint *, 4> constraints;
   SmallVector<TypeVariableType *, 2> referencedVars;
   referencedVars.append(extraTypeVars.begin(), extraTypeVars.end());
@@ -335,7 +334,7 @@ static void createConjunction(ConstraintSystem &cs,
     isIsolated = true;
   }
 
-  TypeVarRefCollector paramCollector(cs);
+  TypeVarRefCollector paramCollector(cs, dc, locator);
 
   // Whether we're doing completion, and the conjunction is for a
   // SingleValueStmtExpr, or one of its braces.
@@ -520,7 +519,8 @@ class SyntacticElementConstraintGenerator
   void createConjunction(ArrayRef<ElementInfo> elements,
                          ConstraintLocator *locator, bool isIsolated = false,
                          ArrayRef<TypeVariableType *> extraTypeVars = {}) {
-    ::createConjunction(cs, elements, locator, isIsolated, extraTypeVars);
+    ::createConjunction(cs, context.getAsDeclContext(), elements, locator,
+                        isIsolated, extraTypeVars);
   }
 
   void visitExprPattern(ExprPattern *EP) {
@@ -1520,6 +1520,9 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
 
 void ConstraintSystem::generateConstraints(ArrayRef<ExprPattern *> exprPatterns,
                                            ConstraintLocatorBuilder locator) {
+  assert(!exprPatterns.empty());
+  auto *DC = exprPatterns.front()->getDeclContext();
+
   // Form a conjunction of ExprPattern elements, isolated from the rest of the
   // pattern.
   SmallVector<ElementInfo> elements;
@@ -1532,7 +1535,7 @@ void ConstraintSystem::generateConstraints(ArrayRef<ExprPattern *> exprPatterns,
     elements.push_back(makeElement(EP, getConstraintLocator(EP), context));
   }
   auto *loc = getConstraintLocator(locator);
-  createConjunction(*this, elements, loc, /*isIsolated*/ true,
+  createConjunction(*this, DC, elements, loc, /*isIsolated*/ true,
                     referencedTypeVars);
 }
 

test/Constraints/rdar114402042.swift

@@ -0,0 +1,32 @@
+// RUN: %target-typecheck-verify-swift
+
+// rdar://114402042 - Make sure we connect the SingleValueStmtExpr to the outer
+// closure's return type.
+func foo<T>(_: () -> T) {}
+func bar<T>(_ x: T) {}
+func test() {
+  foo {
+    bar(if true { return } else { return })
+    // expected-error@-1 {{'if' may only be used as expression in return, throw, or as the source of an assignment}}
+    // expected-error@-2 2{{cannot 'return' in 'if' when used as expression}}
+  }
+  foo {
+    bar(if true { { return } } else { { return } })
+    // expected-error@-1 {{'if' may only be used as expression in return, throw, or as the source of an assignment}}
+  }
+}
+
+func baz() -> String? {
+  nil
+}
+
+var x: Int? = {
+  print( // expected-error {{cannot convert value of type '()' to closure result type 'Int?'}}
+    // expected-note@-1 {{to match this opening '('}}
+    switch baz() {
+    case ""?:
+        return nil
+    default:
+        return nil
+    }
+}() // expected-error {{expected ')' in expression list}}