Lift replaceInferableTypes call into getContextualType

Author: Jumhyn

include/swift/Sema/ConstraintSystem.h

@@ -908,6 +908,7 @@ using OpenedTypeMap =
 /// within a constraint system.
 struct ContextualTypeInfo {
   TypeLoc typeLoc;
+
   ContextualTypePurpose purpose;
 
   ContextualTypeInfo() : typeLoc(TypeLoc()), purpose(CTP_Unused) {}
@@ -2343,8 +2344,10 @@ class ConstraintSystem {
     solutionApplicationTargets;
 
   /// Contextual type information for expressions that are part of this
-  /// constraint system.
-  llvm::MapVector<ASTNode, ContextualTypeInfo> contextualTypes;
+  /// constraint system. The second type, if valid, contains the type as it
+  /// should appear in actual constraint. This will have unbound generic types
+  /// opened, placeholder types converted to type variables, etc.
+  llvm::MapVector<ASTNode, std::pair<ContextualTypeInfo, Type>> contextualTypes;
 
   /// Information about each case label item tracked by the constraint system.
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
@@ -3184,21 +3187,47 @@ class ConstraintSystem {
     assert(bool(node) && "Expected non-null expression!");
     assert(contextualTypes.count(node) == 0 &&
            "Already set this contextual type");
-    contextualTypes[node] = {T, purpose};
+    contextualTypes[node] = {{T, purpose}, Type()};
   }
 
   Optional<ContextualTypeInfo> getContextualTypeInfo(ASTNode node) const {
     auto known = contextualTypes.find(node);
     if (known == contextualTypes.end())
       return None;
-    return known->second;
-  }
-
-  Type getContextualType(ASTNode node) const {
-    auto result = getContextualTypeInfo(node);
-    if (result)
-      return result->typeLoc.getType();
-    return Type();
+    return known->second.first;
+  }
+
+  /// Gets the contextual type recorded for an AST node. When fetching a type
+  /// for use in constraint solving, \c forConstraint should be set to \c true,
+  /// which will ensure that unbound generics have been opened and placeholder
+  /// types have been converted to type variables, etc.
+  Type getContextualType(ASTNode node, bool forConstraint = false) {
+    if (forConstraint) {
+      auto known = contextualTypes.find(node);
+      if (known == contextualTypes.end())
+        return Type();
+
+      // If we've already computed a type for use in the constraint system,
+      // use that.
+      if (known->second.second)
+        return known->second.second;
+
+      // Otherwise, compute a type that can be used in a constraint and record
+      // it.
+      auto info = known->second.first;
+
+      auto *locator = getConstraintLocator(
+          node, LocatorPathElt::ContextualType(info.purpose));
+      known->second.second = replaceInferableTypesWithTypeVars(info.getType(),
+                                                               locator);
+
+      return known->second.second;
+    } else {
+      auto result = getContextualTypeInfo(node);
+      if (result)
+        return result->getType();
+      return Type();
+    }
   }
 
   TypeLoc getContextualTypeLoc(ASTNode node) const {

lib/Sema/CSGen.cpp

@@ -1732,8 +1732,9 @@ namespace {
 
       // If a contextual type exists for this expression, apply it directly.
       if (contextualType && ConstraintSystem::isArrayType(contextualType)) {
-        contextualType = CS.replaceInferableTypesWithTypeVars(
-            contextualType, CS.getConstraintLocator(expr));
+        // Now that we know we're actually going to use the type, get the
+        // version for use in a constraint.
+        contextualType = CS.getContextualType(expr, /*forConstraint=*/true);
         Optional<Type> arrayElementType =
             ConstraintSystem::isArrayType(contextualType);
         CS.addConstraint(ConstraintKind::LiteralConformsTo, contextualType,
@@ -1838,15 +1839,18 @@ namespace {
       auto locator = CS.getConstraintLocator(expr);
       auto contextualType = CS.getContextualType(expr);
       auto contextualPurpose = CS.getContextualTypePurpose(expr);
-      auto openedType =
-          CS.openOpaqueType(contextualType, contextualPurpose, locator);
 
       // If a contextual type exists for this expression and is a dictionary
       // type, apply it directly.
-      if (openedType && ConstraintSystem::isDictionaryType(openedType)) {
+      if (contextualType && ConstraintSystem::isDictionaryType(contextualType)) {
+        // Now that we know we're actually going to use the type, get the
+        // version for use in a constraint.
+        contextualType = CS.getContextualType(expr, /*forConstraint=*/true);
+        auto openedType =
+            CS.openOpaqueType(contextualType, contextualPurpose, locator);
         openedType = CS.replaceInferableTypesWithTypeVars(
             openedType, CS.getConstraintLocator(expr));
-        Optional<std::pair<Type, Type>> dictionaryKeyValue =
+        auto dictionaryKeyValue =
             ConstraintSystem::isDictionaryType(openedType);
         Type contextualDictionaryKeyType;
         Type contextualDictionaryValueType;

lib/Sema/CSSolver.cpp

@@ -180,8 +180,9 @@ Solution ConstraintSystem::finalize() {
   }
 
   // Remember contextual types.
-  solution.contextualTypes.assign(
-      contextualTypes.begin(), contextualTypes.end());
+  for (auto &entry : contextualTypes) {
+    solution.contextualTypes.push_back({entry.first, entry.second.first});
+  }
 
   solution.solutionApplicationTargets = solutionApplicationTargets;
   solution.caseLabelItems = caseLabelItems;

lib/Sema/TypeCheckConstraints.cpp

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