Merge pull request swiftlang#29770 from DougGregor/builder-transform-node-types

[Constraint solver] Fix handling of node types for function builder application

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -1101,6 +1101,9 @@ Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
     // The system was salvaged; continue on as if nothing happened.
   }
 
+  // FIXME: Shouldn't need to do this.
+  cs.applySolution(solutions.front());
+
   // Apply the solution to the function body.
   if (auto result = cs.applySolution(
           solutions.front(),

lib/Sema/CSApply.cpp

@@ -7538,14 +7538,8 @@ ProtocolConformanceRef Solution::resolveConformance(
 }
 
 Type Solution::getType(const Expr *expr) const {
-  auto result = llvm::find_if(
-      addedNodeTypes, [&](const std::pair<TypedNode, Type> &node) -> bool {
-        if (auto *e = node.first.dyn_cast<const Expr *>())
-          return expr == e;
-        return false;
-      });
-
-  if (result != addedNodeTypes.end())
+  auto result = nodeTypes.find(expr);
+  if (result != nodeTypes.end())
     return result->second;
 
   auto &cs = getConstraintSystem();

lib/Sema/CSSolver.cpp

@@ -164,8 +164,8 @@ Solution ConstraintSystem::finalize() {
   solution.DefaultedConstraints.insert(DefaultedConstraints.begin(),
                                        DefaultedConstraints.end());
 
-  for (auto &nodeType : addedNodeTypes) {
-    solution.addedNodeTypes.insert(nodeType);
+  for (auto &nodeType : NodeTypes) {
+    solution.nodeTypes.insert(nodeType);
   }
 
   // Remember contextual types.
@@ -231,9 +231,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
                               solution.DefaultedConstraints.end());
 
   // Add the node types back.
-  for (auto &nodeType : solution.addedNodeTypes) {
-    if (!hasType(nodeType.first))
-      setType(nodeType.first, nodeType.second);
+  for (auto &nodeType : solution.nodeTypes) {
+    setType(nodeType.first, nodeType.second);
   }
 
   // Add the contextual types.
@@ -511,8 +510,13 @@ ConstraintSystem::SolverScope::~SolverScope() {
   truncate(cs.DefaultedConstraints, numDefaultedConstraints);
 
   // Remove any node types we registered.
-  for (unsigned i : range(numAddedNodeTypes, cs.addedNodeTypes.size())) {
-    cs.eraseType(cs.addedNodeTypes[i].first);
+  for (unsigned i :
+           reverse(range(numAddedNodeTypes, cs.addedNodeTypes.size()))) {
+    TypedNode node = cs.addedNodeTypes[i].first;
+    if (Type oldType = cs.addedNodeTypes[i].second)
+      cs.NodeTypes[node] = oldType;
+    else
+      cs.NodeTypes.erase(node);
   }
   truncate(cs.addedNodeTypes, numAddedNodeTypes);
 

lib/Sema/ConstraintSystem.h

@@ -856,7 +856,7 @@ class Solution {
   llvm::SmallPtrSet<ConstraintLocator *, 2> DefaultedConstraints;
 
   /// The node -> type mappings introduced by this solution.
-  llvm::MapVector<TypedNode, Type> addedNodeTypes;
+  llvm::MapVector<TypedNode, Type> nodeTypes;
 
   /// Contextual types introduced by this solution.
   std::vector<std::pair<const Expr *, ContextualTypeInfo>> contextualTypes;
@@ -953,8 +953,8 @@ class Solution {
 
   /// Retrieve the type of the given node, as recorded in this solution.
   Type getType(TypedNode node) const {
-    auto known = addedNodeTypes.find(node);
-    assert(known != addedNodeTypes.end());
+    auto known = nodeTypes.find(node);
+    assert(known != nodeTypes.end());
     return known->second;
   }
 
@@ -1479,15 +1479,12 @@ class ConstraintSystem {
   /// from declared parameters/result and body.
   llvm::MapVector<const ClosureExpr *, FunctionType *> ClosureTypes;
 
-  /// Maps expression types used within all portions of the constraint
-  /// system, instead of directly using the types on the expression
-  /// nodes themselves. This allows us to typecheck an expression and
+  /// Maps node types used within all portions of the constraint
+  /// system, instead of directly using the types on the
+  /// nodes themselves. This allows us to typecheck and
   /// run through various diagnostics passes without actually mutating
-  /// the types on the expression nodes.
-  llvm::DenseMap<const Expr *, TypeBase *> ExprTypes;
-  llvm::DenseMap<const TypeLoc *, TypeBase *> TypeLocTypes;
-  llvm::DenseMap<const VarDecl *, TypeBase *> VarTypes;
-  llvm::DenseMap<const Pattern *, TypeBase *> PatternTypes;
+  /// the types on the nodes.
+  llvm::MapVector<TypedNode, Type> NodeTypes;
   llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, TypeBase *>
       KeyPathComponentTypes;
 
@@ -1555,7 +1552,8 @@ class ConstraintSystem {
   SmallVector<std::pair<ConstraintLocator *, OpenedArchetypeType *>, 4>
     OpenedExistentialTypes;
 
-  /// The node -> type mappings introduced by generating constraints.
+  /// The nodes for which we have produced types, along with the prior type
+  /// each node had before introducing this type.
   llvm::SmallVector<std::pair<TypedNode, Type>, 8> addedNodeTypes;
 
   std::vector<std::pair<ConstraintLocator *, ProtocolConformanceRef>>
@@ -2238,19 +2236,12 @@ class ConstraintSystem {
     assert(type && "Expected non-null type");
 
     // Record the type.
-    if (auto expr = node.dyn_cast<const Expr *>()) {
-      ExprTypes[expr] = type.getPointer();
-    } else if (auto typeLoc = node.dyn_cast<const TypeLoc *>()) {
-      TypeLocTypes[typeLoc] = type.getPointer();
-    } else if (auto var = node.dyn_cast<const VarDecl *>()) {
-      VarTypes[var] = type.getPointer();
-    } else {
-      auto pattern = node.get<const Pattern *>();
-      PatternTypes[pattern] = type.getPointer();
-    }
+    Type &entry = NodeTypes[node];
+    Type oldType = entry;
+    entry = type;
 
     // Record the fact that we ascribed a type to this node.
-    addedNodeTypes.push_back({node, type});
+    addedNodeTypes.push_back({node, oldType});
   }
 
   /// Set the type in our type map for a given expression. The side
@@ -2261,49 +2252,20 @@ class ConstraintSystem {
     setType(TypedNode(&L), T);
   }
 
-  /// Erase the type for the given node.
-  void eraseType(TypedNode node) {
-    if (auto expr = node.dyn_cast<const Expr *>()) {
-      ExprTypes.erase(expr);
-    } else if (auto typeLoc = node.dyn_cast<const TypeLoc *>()) {
-      TypeLocTypes.erase(typeLoc);
-    } else if (auto var = node.dyn_cast<const VarDecl *>()) {
-      VarTypes.erase(var);
-    } else {
-      auto pattern = node.get<const Pattern *>();
-      PatternTypes.erase(pattern);
-    }
-  }
-
   void setType(KeyPathExpr *KP, unsigned I, Type T) {
     assert(KP && "Expected non-null key path parameter!");
     assert(T && "Expected non-null type!");
     KeyPathComponentTypes[std::make_pair(KP, I)] = T.getPointer();
   }
 
-  /// Check to see if we have a type for an expression.
-  bool hasType(const Expr *E) const {
-    assert(E != nullptr && "Expected non-null expression!");
-    return ExprTypes.find(E) != ExprTypes.end();
-  }
-
   bool hasType(const TypeLoc &L) const {
     return hasType(TypedNode(&L));
   }
 
   /// Check to see if we have a type for a node.
   bool hasType(TypedNode node) const {
     assert(!node.isNull() && "Expected non-null node");
-    if (auto expr = node.dyn_cast<const Expr *>()) {
-      return ExprTypes.find(expr) != ExprTypes.end();
-    } else if (auto typeLoc = node.dyn_cast<const TypeLoc *>()) {
-      return TypeLocTypes.find(typeLoc) != TypeLocTypes.end();
-    } else if (auto var = node.dyn_cast<const VarDecl *>()) {
-      return VarTypes.find(var) != VarTypes.end();
-    } else {
-      auto pattern = node.get<const Pattern *>();
-      return PatternTypes.find(pattern) != PatternTypes.end();
-    }
+    return NodeTypes.count(node) > 0;
   }
 
   bool hasType(const KeyPathExpr *KP, unsigned I) const {
@@ -2312,35 +2274,29 @@ class ConstraintSystem {
               != KeyPathComponentTypes.end();
   }
 
-  /// Get the type for an expression.
-  Type getType(const Expr *E) const {
-    assert(hasType(E) && "Expected type to have been set!");
+  /// Get the type for an node.
+  Type getType(TypedNode 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 ExprTypes.find(E)->second;
+    return NodeTypes.find(node)->second;
   }
 
   Type getType(const TypeLoc &L) const {
-    assert(hasType(L) && "Expected type to have been set!");
-    return TypeLocTypes.find(&L)->second;
-  }
-
-  Type getType(const VarDecl *VD) const {
-    assert(hasType(VD) && "Expected type to have been set!");
-    return VarTypes.find(VD)->second;
+    return getType(TypedNode(&L));
   }
 
   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;
   }
 
-  /// Retrieve the type of the variable, if known.
-  Type getTypeIfAvailable(const VarDecl *VD) const {
-    auto known = VarTypes.find(VD);
-    if (known == VarTypes.end())
+  /// Retrieve the type of the node, if known.
+  Type getTypeIfAvailable(TypedNode node) const {
+    auto known = NodeTypes.find(node);
+    if (known == NodeTypes.end())
       return Type();
 
     return known->second;

test/Constraints/function_builder_diags.swift

@@ -327,3 +327,17 @@ func checkImplicitSelfInClosure() {
     }
   }
 }
+
+// rdar://problem/59239224 - crash because some nodes don't have type
+// information during solution application.
+struct X<T> {
+  init(_: T) { }
+}
+
+@TupleBuilder func foo(cond: Bool) -> some Any {
+  if cond {
+    tuplify(cond) { x in
+      X(x)
+    }
+  }
+}