[Constraint system] Save contextual types into each Solution

We’ll want to pull these back out for diagnostic purposes.

Author: DougGregor

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;