[ConstraintSystem] Switch getContextualType and its variants to use TypedNode

Author: xedin

lib/Sema/CSDiagnostics.h

@@ -119,17 +119,17 @@ class FailureDiagnostic {
 
   Type getContextualType(TypedNode anchor) const {
     auto &cs = getConstraintSystem();
-    return cs.getContextualType(anchor.get<const Expr *>());
+    return cs.getContextualType(anchor);
   }
 
   TypeLoc getContextualTypeLoc(TypedNode anchor) const {
     auto &cs = getConstraintSystem();
-    return cs.getContextualTypeLoc(anchor.get<const Expr *>());
+    return cs.getContextualTypeLoc(anchor);
   }
 
   ContextualTypePurpose getContextualTypePurpose(TypedNode anchor) const {
     auto &cs = getConstraintSystem();
-    return cs.getContextualTypePurpose(anchor.get<const Expr *>());
+    return cs.getContextualTypePurpose(anchor);
   }
 
   DeclContext *getDC() const {

lib/Sema/CSFix.cpp

@@ -189,8 +189,7 @@ bool MissingConformance::diagnose(const Solution &solution, bool asNote) const {
 
   if (IsContextual) {
     auto &cs = solution.getConstraintSystem();
-    auto context =
-        cs.getContextualTypePurpose(locator->getAnchor().get<const Expr *>());
+    auto context = cs.getContextualTypePurpose(locator->getAnchor());
     MissingContextualConformanceFailure failure(
         solution, context, NonConformingType, ProtocolType, locator);
     return failure.diagnose(asNote);
@@ -265,7 +264,7 @@ getStructuralTypeContext(const Solution &solution, ConstraintLocator *locator) {
            locator->isLastElement<LocatorPathElt::FunctionArgument>());
 
     auto &cs = solution.getConstraintSystem();
-    auto *anchor = locator->getAnchor().get<const Expr *>();
+    auto anchor = locator->getAnchor();
     auto contextualType = cs.getContextualType(anchor);
     auto exprType = cs.getType(anchor);
     return std::make_tuple(cs.getContextualTypePurpose(anchor), exprType,
@@ -314,8 +313,7 @@ bool AllowTupleTypeMismatch::coalesceAndDiagnose(
   Type toType;
 
   if (getFromType()->is<TupleType>() && getToType()->is<TupleType>()) {
-    purpose =
-        cs.getContextualTypePurpose(locator->getAnchor().get<const Expr *>());
+    purpose = cs.getContextualTypePurpose(locator->getAnchor());
     fromType = getFromType();
     toType = getToType();
   } else if (auto contextualTypeInfo =

lib/Sema/ConstraintSystem.h

@@ -917,7 +917,7 @@ class Solution {
   llvm::MapVector<TypedNode, Type> nodeTypes;
 
   /// Contextual types introduced by this solution.
-  std::vector<std::pair<const Expr *, ContextualTypeInfo>> contextualTypes;
+  std::vector<std::pair<TypedNode, ContextualTypeInfo>> contextualTypes;
 
   /// Maps AST nodes to their solution application targets.
   llvm::MapVector<SolutionApplicationTargetsKey, SolutionApplicationTarget>
@@ -1719,7 +1719,7 @@ class ConstraintSystem {
 
   /// Contextual type information for expressions that are part of this
   /// constraint system.
-  llvm::MapVector<const Expr *, ContextualTypeInfo> contextualTypes;
+  llvm::MapVector<TypedNode, ContextualTypeInfo> contextualTypes;
 
   /// Information about each case label item tracked by the constraint system.
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
@@ -2527,37 +2527,37 @@ class ConstraintSystem {
     return E;
   }
 
-  void setContextualType(
-      const Expr *expr, TypeLoc T, ContextualTypePurpose purpose) {
-    assert(expr != nullptr && "Expected non-null expression!");
-    assert(contextualTypes.count(expr) == 0 &&
+  void setContextualType(TypedNode node, TypeLoc T,
+                         ContextualTypePurpose purpose) {
+    assert(bool(node) && "Expected non-null expression!");
+    assert(contextualTypes.count(node) == 0 &&
            "Already set this contextual type");
-    contextualTypes[expr] = { T, purpose };
+    contextualTypes[node] = {T, purpose};
   }
 
-  Optional<ContextualTypeInfo> getContextualTypeInfo(const Expr *expr) const {
-    auto known = contextualTypes.find(expr);
+  Optional<ContextualTypeInfo> getContextualTypeInfo(TypedNode node) const {
+    auto known = contextualTypes.find(node);
     if (known == contextualTypes.end())
       return None;
     return known->second;
   }
 
-  Type getContextualType(const Expr *expr) const {
-    auto result = getContextualTypeInfo(expr);
+  Type getContextualType(TypedNode node) const {
+    auto result = getContextualTypeInfo(node);
     if (result)
       return result->typeLoc.getType();
     return Type();
   }
 
-  TypeLoc getContextualTypeLoc(const Expr *expr) const {
-    auto result = getContextualTypeInfo(expr);
+  TypeLoc getContextualTypeLoc(TypedNode node) const {
+    auto result = getContextualTypeInfo(node);
     if (result)
       return result->typeLoc;
     return TypeLoc();
   }
 
-  ContextualTypePurpose getContextualTypePurpose(const Expr *expr) const {
-    auto result = getContextualTypeInfo(expr);
+  ContextualTypePurpose getContextualTypePurpose(TypedNode node) const {
+    auto result = getContextualTypeInfo(node);
     if (result)
       return result->purpose;
     return CTP_Unused;