[Constraint system] Set the contextual type for condition expressions.

lib/Sema/CSGen.cpp

@@ -3821,6 +3821,7 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
     return true;
   }
 
+  Type boolTy = boolDecl->getDeclaredType();
   for (const auto &condElement : condition) {
     switch (condElement.getKind()) {
     case StmtConditionElement::CK_Availability:
@@ -3829,15 +3830,18 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
 
     case StmtConditionElement::CK_Boolean: {
       Expr *condExpr = condElement.getBoolean();
+      setContextualType(condExpr, TypeLoc::withoutLoc(boolTy), CTP_Condition);
+
       condExpr = generateConstraints(condExpr, dc);
       if (!condExpr) {
         return true;
       }
 
       addConstraint(ConstraintKind::Conversion,
                     getType(condExpr),
-                    boolDecl->getDeclaredType(),
-                    getConstraintLocator(condExpr));
+                    boolTy,
+                    getConstraintLocator(condExpr,
+                                         LocatorPathElt::ContextualType()));
       continue;
     }
 

lib/Sema/CSSolver.cpp

@@ -168,6 +168,10 @@ Solution ConstraintSystem::finalize() {
     solution.addedNodeTypes.insert(nodeType);
   }
 
+  // Remember contextual types.
+  solution.contextualTypes.assign(
+      contextualTypes.begin(), contextualTypes.end());
+
   for (auto &e : CheckedConformances)
     solution.Conformances.push_back({e.first, e.second});
 
@@ -232,6 +236,14 @@ void ConstraintSystem::applySolution(const Solution &solution) {
       setType(nodeType.first, nodeType.second);
   }
 
+  // Add the contextual types.
+  for (const auto &contextualType : solution.contextualTypes) {
+    if (!getContextualTypeInfo(contextualType.first)) {
+      setContextualType(contextualType.first, contextualType.second.typeLoc,
+                        contextualType.second.purpose);
+    }
+  }
+
   // Register the conformances checked along the way to arrive to solution.
   for (auto &conformance : solution.Conformances)
     CheckedConformances.push_back(conformance);

lib/Sema/ConstraintSystem.h

@@ -787,6 +787,15 @@ using OpenedType = std::pair<GenericTypeParamType *, TypeVariableType *>;
 using OpenedTypeMap =
     llvm::DenseMap<GenericTypeParamType *, TypeVariableType *>;
 
+/// Describes contextual type information about a particular expression
+/// within a constraint system.
+struct ContextualTypeInfo  {
+  TypeLoc typeLoc;
+  ContextualTypePurpose purpose;
+
+  Type getType() const { return typeLoc.getType(); }
+};
+
 /// A complete solution to a constraint system.
 ///
 /// A solution to a constraint system consists of type variable bindings to
@@ -849,6 +858,9 @@ class Solution {
   /// The node -> type mappings introduced by this solution.
   llvm::MapVector<TypedNode, Type> addedNodeTypes;
 
+  /// Contextual types introduced by this solution.
+  std::vector<std::pair<const Expr *, ContextualTypeInfo>> contextualTypes;
+
   std::vector<std::pair<ConstraintLocator *, ProtocolConformanceRef>>
       Conformances;
 
@@ -1300,13 +1312,6 @@ class ConstraintSystem {
   llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, TypeBase *>
       KeyPathComponentTypes;
 
-  struct ContextualTypeInfo  {
-    TypeLoc typeLoc;
-    ContextualTypePurpose purpose;
-
-    Type getType() const { return typeLoc.getType(); }
-  };
-
   /// Contextual type information for expressions that are part of this
   /// constraint system.
   llvm::MapVector<const Expr *, ContextualTypeInfo> contextualTypes;
@@ -2174,35 +2179,36 @@ class ConstraintSystem {
     return E;
   }
 
-  void setContextualType(Expr *expr, TypeLoc T, ContextualTypePurpose purpose) {
+  void setContextualType(
+      const Expr *expr, TypeLoc T, ContextualTypePurpose purpose) {
     assert(expr != nullptr && "Expected non-null expression!");
     assert(contextualTypes.count(expr) == 0 &&
            "Already set this contextual type");
     contextualTypes[expr] = { T, purpose };
   }
 
-  Optional<ContextualTypeInfo> getContextualTypeInfo(Expr *expr) const {
+  Optional<ContextualTypeInfo> getContextualTypeInfo(const Expr *expr) const {
     auto known = contextualTypes.find(expr);
     if (known == contextualTypes.end())
       return None;
     return known->second;
   }
 
-  Type getContextualType(Expr *expr) const {
+  Type getContextualType(const Expr *expr) const {
     auto result = getContextualTypeInfo(expr);
     if (result)
       return result->typeLoc.getType();
     return Type();
   }
 
-  TypeLoc getContextualTypeLoc(Expr *expr) const {
+  TypeLoc getContextualTypeLoc(const Expr *expr) const {
     auto result = getContextualTypeInfo(expr);
     if (result)
       return result->typeLoc;
     return TypeLoc();
   }
 
-  ContextualTypePurpose getContextualTypePurpose(Expr *expr) const {
+  ContextualTypePurpose getContextualTypePurpose(const Expr *expr) const {
     auto result = getContextualTypeInfo(expr);
     if (result)
       return result->purpose;

test/Constraints/function_builder_diags.swift

@@ -257,7 +257,7 @@ struct MyTuplifiedStruct {
   }
 }
 
-// Check that we're performing syntactic use diagnostics/
+// Check that we're performing syntactic use diagnostics.
 func acceptMetatype<T>(_: T.Type) -> Bool { true }
 
 func syntacticUses<T>(_: T) {
@@ -269,3 +269,18 @@ func syntacticUses<T>(_: T) {
     }
   }
 }
+
+// Check custom diagnostics within "if" conditions.
+struct HasProperty {
+  var property: Bool = false
+}
+
+func checkConditions(cond: Bool) {
+  var x = HasProperty()
+
+  tuplify(cond) { value in
+    if x.property = value { // expected-error{{use of '=' in a boolean context, did you mean '=='?}}
+      "matched it"
+    }
+  }
+}