Merge pull request #68331 from hamishknight/linked-in

hamishknight · web-flow · commit 43b6117d28c8 · 2023-09-06T11:05:36.000+01:00
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -6238,18 +6238,29 @@ class ConjunctionElementProducer : public BindingProducer<ConjunctionElement> {
   }
 };
 
-/// Find any references to not yet resolved outer VarDecls (including closure
-/// parameters) used in the body of a conjunction element (e.g closures, taps,
-/// if/switch expressions). 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 VarRefCollector : public ASTWalker {
+/// Find any references to external type variables used in the body of a
+/// conjunction element (e.g closures, taps, if/switch expressions).
+///
+/// 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:
-  VarRefCollector(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);
@@ -6259,6 +6270,8 @@ class VarRefCollector : 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
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -866,7 +866,7 @@ namespace {
   };
 } // end anonymous namespace
 
-void VarRefCollector::inferTypeVars(Decl *D) {
+void TypeVarRefCollector::inferTypeVars(Decl *D) {
   // We're only interested in VarDecls.
   if (!isa_and_nonnull<VarDecl>(D))
     return;
@@ -881,7 +881,10 @@ void VarRefCollector::inferTypeVars(Decl *D) {
 }
 
 ASTWalker::PreWalkResult<Expr *>
-VarRefCollector::walkToExprPre(Expr *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 @@ VarRefCollector::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()) {
-            VarRefCollector 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);
 
-      VarRefCollector 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 =
diff --git a/lib/Sema/CSSyntacticElement.cpp b/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;
   }
 
-  VarRefCollector paramCollector(cs);
+  TypeVarRefCollector paramCollector(cs, dc, locator);
 
   // Whether we're doing completion, and the conjunction is for a
   // SingleValueStmtExpr, or one of its braces.
@@ -517,6 +516,13 @@ class SyntacticElementConstraintGenerator
                                       ConstraintLocator *locator)
       : cs(cs), context(context), locator(locator) {}
 
+  void createConjunction(ArrayRef<ElementInfo> elements,
+                         ConstraintLocator *locator, bool isIsolated = false,
+                         ArrayRef<TypeVariableType *> extraTypeVars = {}) {
+    ::createConjunction(cs, context.getAsDeclContext(), elements, locator,
+                        isIsolated, extraTypeVars);
+  }
+
   void visitExprPattern(ExprPattern *EP) {
     auto target = SyntacticElementTarget::forExprPattern(EP);
 
@@ -859,7 +865,7 @@ class SyntacticElementConstraintGenerator
     if (auto *join = context.ElementJoin.getPtrOrNull())
       elements.push_back(makeJoinElement(cs, join, locator));
 
-    createConjunction(cs, elements, locator);
+    createConjunction(elements, locator);
   }
 
   void visitGuardStmt(GuardStmt *guardStmt) {
@@ -868,7 +874,7 @@ class SyntacticElementConstraintGenerator
     visitStmtCondition(guardStmt, elements, locator);
     elements.push_back(makeElement(guardStmt->getBody(), locator));
 
-    createConjunction(cs, elements, locator);
+    createConjunction(elements, locator);
   }
 
   void visitWhileStmt(WhileStmt *whileStmt) {
@@ -877,16 +883,15 @@ class SyntacticElementConstraintGenerator
     visitStmtCondition(whileStmt, elements, locator);
     elements.push_back(makeElement(whileStmt->getBody(), locator));
 
-    createConjunction(cs, elements, locator);
+    createConjunction(elements, locator);
   }
 
   void visitDoStmt(DoStmt *doStmt) {
     visitBraceStmt(doStmt->getBody());
   }
 
   void visitRepeatWhileStmt(RepeatWhileStmt *repeatWhileStmt) {
-    createConjunction(cs,
-                      {makeElement(repeatWhileStmt->getCond(),
+    createConjunction({makeElement(repeatWhileStmt->getCond(),
                                    cs.getConstraintLocator(
                                        locator, ConstraintLocator::Condition),
                                    getContextForCondition()),
@@ -895,8 +900,7 @@ class SyntacticElementConstraintGenerator
   }
 
   void visitPoundAssertStmt(PoundAssertStmt *poundAssertStmt) {
-    createConjunction(cs,
-                      {makeElement(poundAssertStmt->getCondition(),
+    createConjunction({makeElement(poundAssertStmt->getCondition(),
                                    cs.getConstraintLocator(
                                        locator, ConstraintLocator::Condition),
                                    getContextForCondition())},
@@ -913,12 +917,10 @@ class SyntacticElementConstraintGenerator
     auto *errorExpr = throwStmt->getSubExpr();
 
     createConjunction(
-        cs,
-        {makeElement(
-            errorExpr,
-            cs.getConstraintLocator(
-                locator, LocatorPathElt::SyntacticElement(errorExpr)),
-            {errType, CTP_ThrowStmt})},
+        {makeElement(errorExpr,
+                     cs.getConstraintLocator(
+                         locator, LocatorPathElt::SyntacticElement(errorExpr)),
+                     {errType, CTP_ThrowStmt})},
         locator);
   }
 
@@ -939,12 +941,10 @@ class SyntacticElementConstraintGenerator
     auto *selfExpr = discardStmt->getSubExpr();
 
     createConjunction(
-        cs,
-        {makeElement(
-            selfExpr,
-            cs.getConstraintLocator(
-                locator, LocatorPathElt::SyntacticElement(selfExpr)),
-            {nominalType, CTP_DiscardStmt})},
+        {makeElement(selfExpr,
+                     cs.getConstraintLocator(
+                         locator, LocatorPathElt::SyntacticElement(selfExpr)),
+                     {nominalType, CTP_DiscardStmt})},
         locator);
   }
 
@@ -962,7 +962,7 @@ class SyntacticElementConstraintGenerator
     // Body of the `for-in` loop.
     elements.push_back(makeElement(forEachStmt->getBody(), stmtLoc));
 
-    createConjunction(cs, elements, locator);
+    createConjunction(elements, locator);
   }
 
   void visitSwitchStmt(SwitchStmt *switchStmt) {
@@ -990,7 +990,7 @@ class SyntacticElementConstraintGenerator
     if (auto *join = context.ElementJoin.getPtrOrNull())
       elements.push_back(makeJoinElement(cs, join, switchLoc));
 
-    createConjunction(cs, elements, switchLoc);
+    createConjunction(elements, switchLoc);
   }
 
   void visitDoCatchStmt(DoCatchStmt *doStmt) {
@@ -1007,7 +1007,7 @@ class SyntacticElementConstraintGenerator
     for (auto *catchStmt : doStmt->getCatches())
       elements.push_back(makeElement(catchStmt, doLoc));
 
-    createConjunction(cs, elements, doLoc);
+    createConjunction(elements, doLoc);
   }
 
   void visitCaseStmt(CaseStmt *caseStmt) {
@@ -1040,7 +1040,7 @@ class SyntacticElementConstraintGenerator
 
     elements.push_back(makeElement(caseStmt->getBody(), caseLoc));
 
-    createConjunction(cs, elements, caseLoc);
+    createConjunction(elements, caseLoc);
   }
 
   void visitBraceStmt(BraceStmt *braceStmt) {
@@ -1127,7 +1127,7 @@ class SyntacticElementConstraintGenerator
         // want to type-check captures to make sure that the context
         // is valid.
         if (captureList)
-          createConjunction(cs, elements, locator);
+          createConjunction(elements, locator);
 
         return;
       }
@@ -1195,7 +1195,7 @@ class SyntacticElementConstraintGenerator
           contextInfo.value_or(ContextualTypeInfo()), isDiscarded));
     }
 
-    createConjunction(cs, elements, locator);
+    createConjunction(elements, locator);
   }
 
   void visitReturnStmt(ReturnStmt *returnStmt) {
@@ -1282,7 +1282,7 @@ class SyntacticElementConstraintGenerator
     auto resultElt = makeElement(resultExpr, locator,
                                  contextInfo.value_or(ContextualTypeInfo()),
                                  /*isDiscarded=*/false);
-    createConjunction(cs, {resultElt}, locator);
+    createConjunction({resultElt}, locator);
   }
 
   ContextualTypeInfo getContextualResultInfo() const {
@@ -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);
 }
 
diff --git a/test/Constraints/rdar114402042.swift b/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}}
