[CS] Solve all conjunctions in source order

include/swift/AST/Expr.h

@@ -880,12 +880,8 @@ class TapExpr : public Expr {
   void setBody(BraceStmt * b) { Body = b; }
 
   SourceLoc getLoc() const { return SubExpr ? SubExpr->getLoc() : SourceLoc(); }
-
-  SourceLoc getStartLoc() const {
-    return SubExpr ? SubExpr->getStartLoc() : SourceLoc();
-  }
 
-  SourceLoc getEndLoc() const;
+  SourceRange getSourceRange() const;
 
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Tap;

lib/AST/Expr.cpp

@@ -2611,29 +2611,33 @@ void InterpolatedStringLiteralExpr::forEachSegment(ASTContext &Ctx,
 TapExpr::TapExpr(Expr * SubExpr, BraceStmt *Body)
     : Expr(ExprKind::Tap, /*Implicit=*/true),
       SubExpr(SubExpr), Body(Body) {
-  if (Body) {
-    assert(!Body->empty() &&
-         Body->getFirstElement().isDecl(DeclKind::Var) &&
-         "First element of Body should be a variable to init with the subExpr");
-  }
+  assert(Body);
+  assert(!Body->empty() &&
+       Body->getFirstElement().isDecl(DeclKind::Var) &&
+       "First element of Body should be a variable to init with the subExpr");
 }
 
 VarDecl * TapExpr::getVar() const {
   return dyn_cast<VarDecl>(Body->getFirstElement().dyn_cast<Decl *>());
 }
 
-SourceLoc TapExpr::getEndLoc() const {
-  // Include the body in the range, assuming the body follows the SubExpr.
-  // Also, be (perhaps overly) defensive about null pointers & invalid
-  // locations.
-  if (auto *const b = getBody()) {
-    const auto be = b->getEndLoc();
-    if (be.isValid())
-      return be;
-  }
-  if (auto *const se = getSubExpr())
-    return se->getEndLoc();
-  return SourceLoc();
+SourceRange TapExpr::getSourceRange() const {
+  if (!SubExpr)
+    return Body->getSourceRange();
+
+  SourceLoc start = SubExpr->getStartLoc();
+  if (!start.isValid())
+    start = Body->getStartLoc();
+  if (!start.isValid())
+    return SourceRange();
+
+  SourceLoc end = Body->getEndLoc();
+  if (!end.isValid())
+    end = SubExpr->getEndLoc();
+  if (!end.isValid())
+    return SourceRange();
+
+  return SourceRange(start, end);
 }
 
 RegexLiteralExpr *

lib/Sema/CSSolver.cpp

@@ -2438,25 +2438,31 @@ Constraint *ConstraintSystem::selectConjunction() {
 
   auto &SM = getASTContext().SourceMgr;
 
-  // All of the multi-statement closures should be solved in order of their
-  // apperance in the source.
-  llvm::sort(
-      conjunctions, [&](Constraint *conjunctionA, Constraint *conjunctionB) {
+  // Conjunctions should be solved in order of their apperance in the source.
+  // This is important because once a conjunction is solved, we don't re-visit
+  // it, so we need to make sure we don't solve it before another conjuntion
+  // that could provide it with necessary type information. Source order
+  // provides an easy to reason about and quick way of establishing this.
+  return *std::min_element(
+      conjunctions.begin(), conjunctions.end(),
+      [&](Constraint *conjunctionA, Constraint *conjunctionB) {
         auto *locA = conjunctionA->getLocator();
         auto *locB = conjunctionB->getLocator();
-
         if (!(locA && locB))
           return false;
 
-        auto *closureA = getAsExpr<ClosureExpr>(locA->getAnchor());
-        auto *closureB = getAsExpr<ClosureExpr>(locB->getAnchor());
+        auto anchorA = locA->getAnchor();
+        auto anchorB = locB->getAnchor();
+        if (!(anchorA && anchorB))
+          return false;
 
-        return closureA && closureB
-                   ? SM.isBeforeInBuffer(closureA->getLoc(), closureB->getLoc())
-                   : false;
-      });
+        auto slocA = anchorA.getStartLoc();
+        auto slocB = anchorB.getStartLoc();
+        if (!(slocA.isValid() && slocB.isValid()))
+          return false;
 
-  return conjunctions.front();
+        return SM.isBeforeInBuffer(slocA, slocB);
+      });
 }
 
 bool DisjunctionChoice::attempt(ConstraintSystem &cs) const {

test/Constraints/rdar113326835.swift

@@ -0,0 +1,46 @@
+// RUN: %target-typecheck-verify-swift
+
+// rdar://113326835 - Make sure we type-check the conjunctions in source order,
+// the first closure should be type-checked before we attempt the
+// TapExpr/SingleValueExpr conjunctions, since they rely on 'k' being resolved.
+
+func global<T>(_ x: T) -> String { "" }
+func global(_ x: Any.Type) -> String { "" }
+
+protocol P {
+  associatedtype X
+}
+
+struct Q<X>: P {
+  init() {}
+  func bar(_: String) -> Self { fatalError() }
+  func qux<U: P>(_: (X) -> U) -> Q<U.X> { fatalError() }
+}
+
+struct J<X>: P {
+  init(_: X) {}
+  func baz<T>(_ transform: (X) -> T) -> Q<T> { fatalError() }
+}
+
+func foo(a: Int) -> Q<String> {
+  J(a)
+    .baz { x in
+      ()
+      return a
+    }
+    .qux { k in
+      Q<String>().bar("\(k)")
+    }
+}
+
+func bar(a: Int) -> Q<String> {
+  J(a)
+    .baz { x in
+      ()
+      return a
+    }
+    .qux { k in
+      Q<String>().bar(if .random() { global(k) } else { global(k) })
+      // expected-error@-1 {{'if' may only be used as expression in return, throw, or as the source of an assignment}}
+    }
+}