[Constraint solver] Allow contextual types for multiple expressions.

Generalize the representation of contextual types so we can store
contextual types for more than one expression. Allow these to be
added/rolled back in solver scopes. Nothing uses this functionality
yet.

Author: DougGregor

lib/Sema/CSSolver.cpp

@@ -441,6 +441,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numFunctionBuilderTransformed = cs.functionBuilderTransformed.size();
   numResolvedOverloads = cs.ResolvedOverloads.size();
   numInferredClosureTypes = cs.ClosureTypes.size();
+  numContextualTypes = cs.contextualTypes.size();
 
   PreviousScore = cs.CurrentScore;
 
@@ -510,6 +511,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any inferred closure types (e.g. used in function builder body).
   truncate(cs.ClosureTypes, numInferredClosureTypes);
 
+  // Remove any contextual types.
+  truncate(cs.contextualTypes, numContextualTypes);
+
   // Reset the previous score.
   cs.CurrentScore = PreviousScore;
 

lib/Sema/ConstraintSystem.h

@@ -1300,17 +1300,21 @@ 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;
+
   /// Maps closure parameters to type variables.
   llvm::DenseMap<const ParamDecl *, TypeVariableType *>
     OpenedParameterTypes;
 
-  /// There can only be a single contextual type on the root of the expression
-  /// being checked.  If specified, this holds its type along with the base
-  /// expression, and the purpose of it.
-  TypeLoc contextualType;
-  Expr *contextualTypeNode = nullptr;
-  ContextualTypePurpose contextualTypePurpose = CTP_Unused;
-  
   /// The set of constraint restrictions used to reach the
   /// current constraint system.
   ///
@@ -1880,6 +1884,9 @@ class ConstraintSystem {
     /// The length of \c ClosureTypes.
     unsigned numInferredClosureTypes;
 
+    /// The length of \c contextualTypes.
+    unsigned numContextualTypes;
+
     /// The previous score.
     Score PreviousScore;
 
@@ -2167,25 +2174,39 @@ class ConstraintSystem {
     return E;
   }
 
-  void setContextualType(Expr *E, TypeLoc T, ContextualTypePurpose purpose) {
-    assert(E != nullptr && "Expected non-null expression!");
-    contextualTypeNode = E;
-    contextualType = T;
-    contextualTypePurpose = purpose;
+  void setContextualType(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 };
   }
 
-  Type getContextualType(Expr *E) const {
-    assert(E != nullptr && "Expected non-null expression!");
-    return E == contextualTypeNode ? contextualType.getType() : Type();
+  Optional<ContextualTypeInfo> getContextualTypeInfo(Expr *expr) const {
+    auto known = contextualTypes.find(expr);
+    if (known == contextualTypes.end())
+      return None;
+    return known->second;
+  }
+
+  Type getContextualType(Expr *expr) const {
+    auto result = getContextualTypeInfo(expr);
+    if (result)
+      return result->typeLoc.getType();
+    return Type();
   }
 
   TypeLoc getContextualTypeLoc(Expr *expr) const {
-    return expr == contextualTypeNode ? contextualType : TypeLoc();
+    auto result = getContextualTypeInfo(expr);
+    if (result)
+      return result->typeLoc;
+    return TypeLoc();
   }
 
   ContextualTypePurpose getContextualTypePurpose(Expr *expr) const {
-    return (expr && expr == contextualTypeNode) ? contextualTypePurpose
-                                                : CTP_Unused;
+    auto result = getContextualTypeInfo(expr);
+    if (result)
+      return result->purpose;
+    return CTP_Unused;
   }
 
   /// Retrieve the constraint locator for the given anchor and

lib/Sema/TypeCheckConstraints.cpp

@@ -3770,9 +3770,9 @@ void ConstraintSystem::print(raw_ostream &out) const {
 
   out << "Score: " << CurrentScore << "\n";
 
-  if (auto ty = contextualType.getType()) {
-    out << "Contextual Type: " << ty.getString(PO);
-    if (TypeRepr *TR = contextualType.getTypeRepr()) {
+  for (const auto &contextualType : contextualTypes) {
+    out << "Contextual Type: " << contextualType.second.getType().getString(PO);
+    if (TypeRepr *TR = contextualType.second.typeLoc.getTypeRepr()) {
       out << " at ";
       TR->getSourceRange().print(out, getASTContext().SourceMgr, /*text*/false);
     }