Sema: Record ASTNode type changes in PreparedOverload

Also, move some methods from ConstraintSystem.h to ConstraintSystem.cpp.

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -3068,7 +3068,7 @@ class ConstraintSystem {
     return nullptr;
   }
 
-  TypeBase* getFavoredType(Expr *E) {
+  TypeBase *getFavoredType(Expr *E) {
     assert(E != nullptr);
     return this->FavoredTypes[E];
   }
@@ -3082,95 +3082,46 @@ class ConstraintSystem {
   ///
   /// The side tables are used through the expression type checker to avoid
   /// mutating nodes until we know we have successfully type-checked them.
-  void setType(ASTNode node, Type type) {
-    ASSERT(!node.isNull() && "Cannot set type information on null node");
-    ASSERT(type && "Expected non-null type");
-
-    // Record the type.
-    Type &entry = NodeTypes[node];
-    Type oldType = entry;
-    entry = type;
-
-    if (oldType.getPointer() != type.getPointer()) {
-      // Record the fact that we assigned a type to this node.
-      if (solverState)
-        recordChange(SolverTrail::Change::RecordedNodeType(node, oldType));
-    }
-  }
+  void setType(ASTNode node, Type type,
+               PreparedOverloadBuilder *preparedOverload=nullptr);
 
   /// Undo the above change.
-  void restoreType(ASTNode node, Type oldType) {
-    ASSERT(!node.isNull() && "Cannot set type information on null node");
-
-    if (oldType) {
-      auto found = NodeTypes.find(node);
-      ASSERT(found != NodeTypes.end());
-      found->second = oldType;
-    } else {
-      bool erased = NodeTypes.erase(node);
-      ASSERT(erased);
-    }
-  }
+  void restoreType(ASTNode node, Type oldType);
 
   /// Check to see if we have a type for a node.
   bool hasType(ASTNode node) const {
-    assert(!node.isNull() && "Expected non-null node");
+    ASSERT(!node.isNull() && "Expected non-null node");
     return NodeTypes.count(node) > 0;
   }
 
   /// Set the type in our type map for a given expression. The side
   /// map is used throughout the expression type checker in order to
   /// avoid mutating expressions until we know we have successfully
   /// type-checked them.
-  void setType(const KeyPathExpr *KP, unsigned I, Type T) {
-    ASSERT(KP && "Expected non-null key path parameter!");
-    ASSERT(T && "Expected non-null type!");
+  void setType(const KeyPathExpr *KP, unsigned I, Type T);
 
-    Type &entry = KeyPathComponentTypes[{KP, I}];
-    Type oldType = entry;
-    entry = T;
-
-    if (oldType.getPointer() != T.getPointer()) {
-      if (solverState) {
-        recordChange(
-          SolverTrail::Change::RecordedKeyPathComponentType(
-            KP, I, oldType));
-      }
-    }
-  }
-
-  void restoreType(const KeyPathExpr *KP, unsigned I, Type T) {
-    ASSERT(KP && "Expected non-null key path parameter!");
-
-    if (T) {
-      auto found = KeyPathComponentTypes.find({KP, I});
-      ASSERT(found != KeyPathComponentTypes.end());
-      found->second = T;
-    } else {
-      bool erased = KeyPathComponentTypes.erase({KP, I});
-      ASSERT(erased);
-    }
-  }
+  void restoreType(const KeyPathExpr *KP, unsigned I, Type T);
 
   bool hasType(const KeyPathExpr *KP, unsigned I) const {
-    assert(KP && "Expected non-null key path parameter!");
+    ASSERT(KP && "Expected non-null key path parameter!");
     return KeyPathComponentTypes.find(std::make_pair(KP, I))
               != KeyPathComponentTypes.end();
   }
 
   /// Get the type for an node.
   Type getType(ASTNode node) const {
-    assert(hasType(node) && "Expected type to have been set!");
-    // FIXME: lvalue differences
-    //    assert((!E->getType() ||
-    //            E->getType()->isEqual(ExprTypes.find(E)->second)) &&
-    //           "Mismatched types!");
-    return NodeTypes.find(node)->second;
+    ASSERT(!node.isNull() && "Expected non-null node");
+    auto found = NodeTypes.find(node);
+    ASSERT(found != NodeTypes.end() && "Expected type to have been set!");
+    return found->second;
   }
 
   Type getType(const KeyPathExpr *KP, unsigned I) const {
-    assert(hasType(KP, I) && "Expected type to have been set!");
-    return KeyPathComponentTypes.find(std::make_pair(KP, I))->second;
+    ASSERT(KP && "Expected non-null key path parameter!");
+    auto found = KeyPathComponentTypes.find(std::make_pair(KP, I));
+    ASSERT(found != KeyPathComponentTypes.end() &&
+           "Expected type to have been set!");
+    return found->second;
   }
 
   /// Retrieve the type of the node, if known.
@@ -3182,11 +3133,6 @@ class ConstraintSystem {
     return known->second;
   }
 
-  /// Retrieve type of the given declaration to be used in
-  /// constraint system, this is better than calling `getType()`
-  /// directly because it accounts of constraint system flags.
-  Type getVarType(const VarDecl *var);
-
   /// Cache the type of the expression argument and return that same
   /// argument.
   template <typename T>

include/swift/Sema/PreparedOverload.h

@@ -88,7 +88,10 @@ struct PreparedOverloadChange {
     AppliedPropertyWrapper,
 
     /// A fix was recorded because a property wrapper application failed.
-    AddedFix
+    AddedFix,
+
+    /// The type of an AST node was changed.
+    RecordedNodeType
   };
 
   /// The kind of change.
@@ -132,7 +135,17 @@ struct PreparedOverloadChange {
       ConstraintFix *TheFix;
       unsigned Impact;
     } Fix;
+
+    /// For ChangeKind::RecordedNodeType.
+    struct {
+      ASTNode Node;
+      Type TheType;
+    } Node;
   };
+
+  PreparedOverloadChange()
+      : Kind(ChangeKind::AddedTypeVariable),
+        TypeVar(nullptr) { }
 };
 
 /// A "pre-cooked" representation of all type variables and constraints
@@ -238,6 +251,14 @@ struct PreparedOverloadBuilder {
     change.Fix.Impact = impact;
     Changes.push_back(change);
   }
+
+  void recordedNodeType(ASTNode node, Type type) {
+    PreparedOverload::Change change;
+    change.Kind = PreparedOverload::Change::RecordedNodeType;
+    change.Node.Node = node;
+    change.Node.TheType = type;
+    Changes.push_back(change);
+  }
 };
 
 }  // end namespace constraints

lib/Sema/CSGen.cpp

@@ -2811,6 +2811,9 @@ namespace {
           // interface type request via `var->getType()` that would
           // attempt to validate `weak` attribute, and produce a
           // diagnostic in the middle of the solver path.
+          //
+          // FIXME: getVarType() is now gone. Is the above still
+          // relevant?
 
           CS.addConstraint(ConstraintKind::OneWayEqual, oneWayVarType,
                            varType, locator);
@@ -5241,7 +5244,8 @@ ConstraintSystem::applyPropertyWrapperToParameter(
                   locator,
                   /*isFavored=*/false,
                   preparedOverload);
-    setType(param->getPropertyWrapperWrappedValueVar(), wrappedValueType);
+    setType(param->getPropertyWrapperWrappedValueVar(), wrappedValueType,
+            preparedOverload);
 
     applyPropertyWrapper(anchor,
                          { wrapperType, PropertyWrapperInitKind::WrappedValue },

lib/Sema/ConstraintSystem.cpp

@@ -905,6 +905,73 @@ ConstraintLocator *ConstraintSystem::getOpenOpaqueLocator(
       { LocatorPathElt::OpenedOpaqueArchetype(opaqueDecl) }, 0);
 }
 
+void ConstraintSystem::setType(ASTNode node, Type type,
+                               PreparedOverloadBuilder *preparedOverload) {
+  ASSERT(PreparingOverload == !!preparedOverload);
+
+  ASSERT(!node.isNull() && "Cannot set type information on null node");
+  ASSERT(type && "Expected non-null type");
+
+  if (preparedOverload) {
+    preparedOverload->recordedNodeType(node, type);
+    return;
+  }
+
+  // Record the type.
+  Type &entry = NodeTypes[node];
+  Type oldType = entry;
+  entry = type;
+
+  if (oldType.getPointer() != type.getPointer()) {
+    // Record the fact that we assigned a type to this node.
+    if (solverState)
+      recordChange(SolverTrail::Change::RecordedNodeType(node, oldType));
+  }
+}
+
+void ConstraintSystem::restoreType(ASTNode node, Type oldType) {
+  ASSERT(!node.isNull() && "Cannot set type information on null node");
+
+  if (oldType) {
+    auto found = NodeTypes.find(node);
+    ASSERT(found != NodeTypes.end());
+    found->second = oldType;
+  } else {
+    bool erased = NodeTypes.erase(node);
+    ASSERT(erased);
+  }
+}
+
+void ConstraintSystem::setType(const KeyPathExpr *KP, unsigned I, Type T) {
+  ASSERT(KP && "Expected non-null key path parameter!");
+  ASSERT(T && "Expected non-null type!");
+
+  Type &entry = KeyPathComponentTypes[{KP, I}];
+  Type oldType = entry;
+  entry = T;
+
+  if (oldType.getPointer() != T.getPointer()) {
+    if (solverState) {
+      recordChange(
+        SolverTrail::Change::RecordedKeyPathComponentType(
+          KP, I, oldType));
+    }
+  }
+}
+
+void ConstraintSystem::restoreType(const KeyPathExpr *KP, unsigned I, Type T) {
+  ASSERT(KP && "Expected non-null key path parameter!");
+
+  if (T) {
+    auto found = KeyPathComponentTypes.find({KP, I});
+    ASSERT(found != KeyPathComponentTypes.end());
+    found->second = T;
+  } else {
+    bool erased = KeyPathComponentTypes.erase({KP, I});
+    ASSERT(erased);
+  }
+}
+
 std::pair<Type, ExistentialArchetypeType *>
 ConstraintSystem::openAnyExistentialType(Type type,
                                          ConstraintLocator *locator) {

lib/Sema/TypeOfReference.cpp

@@ -2834,6 +2834,10 @@ void ConstraintSystem::replayChanges(
     case PreparedOverload::Change::AddedFix:
       recordFix(change.Fix.TheFix, change.Fix.Impact);
       break;
+
+    case PreparedOverload::Change::RecordedNodeType:
+      setType(change.Node.Node, change.Node.TheType);
+      break;
     }
   }
 }