Sema: Recycle ConstraintGraphNode

Author: slavapestov

include/swift/Sema/ConstraintGraph.h

@@ -57,6 +57,13 @@ class ConstraintGraphNode {
   /// Retrieve the type variable this node represents.
   TypeVariableType *getTypeVariable() const { return TypeVar; }
 
+  void reset();
+
+  void initTypeVariable(TypeVariableType *typeVar) {
+    ASSERT(!TypeVar);
+    TypeVar = typeVar;
+  }
+
   /// Retrieve the set of constraints that mention this type variable.
   ///
   /// These are the hyperedges of the graph, connecting this node to
@@ -435,6 +442,9 @@ class ConstraintGraph {
   /// Constraints that are "orphaned" because they contain no type variables.
   SmallVector<Constraint *, 4> OrphanedConstraints;
 
+  /// Unused nodes.
+  SmallVector<ConstraintGraphNode *> FreeList;
+
   /// Increment the number of constraints considered per attempt
   /// to contract constraint graph edges.
   void incrementConstraintsPerContractionCounter();

lib/Sema/ConstraintGraph.cpp

@@ -43,6 +43,9 @@ ConstraintGraph::~ConstraintGraph() {
     delete impl.getGraphNode();
     impl.setGraphNode(nullptr);
   }
+  for (auto *node : FreeList) {
+    delete node;
+  }
 }
 
 #pragma mark Graph accessors
@@ -55,11 +58,20 @@ ConstraintGraph::lookupNode(TypeVariableType *typeVar) {
     assert(impl.getGraphIndex() < TypeVariables.size() && "Out-of-bounds index");
     assert(TypeVariables[impl.getGraphIndex()] == typeVar && 
            "Type variable mismatch");
+    ASSERT(nodePtr->TypeVar == typeVar &&
+           "Use-after-free");
     return { *nodePtr, impl.getGraphIndex() };
   }
 
   // Allocate the new node.
-  auto nodePtr = new ConstraintGraphNode(*this, typeVar);
+  ConstraintGraphNode *nodePtr;
+  if (FreeList.empty())
+    nodePtr = new ConstraintGraphNode(*this, typeVar);
+  else {
+    nodePtr = FreeList.back();
+    FreeList.pop_back();
+    nodePtr->initTypeVariable(typeVar);
+  }
   unsigned index = TypeVariables.size();
   impl.setGraphNode(nodePtr);
   impl.setGraphIndex(index);
@@ -84,6 +96,17 @@ ConstraintGraph::lookupNode(TypeVariableType *typeVar) {
   return { *nodePtr, index };
 }
 
+void ConstraintGraphNode::reset() {
+  ASSERT(TypeVar);
+  TypeVar = nullptr;
+  Bindings.reset();
+  Constraints.clear();
+  ConstraintIndex.clear();
+  ReferencedBy.clear();
+  References.clear();
+  EquivalenceClass.clear();
+}
+
 bool ConstraintGraphNode::forRepresentativeVar() const {
   auto *typeVar = getTypeVariable();
   return typeVar == typeVar->getImpl().getRepresentative(nullptr);
@@ -422,7 +445,9 @@ void ConstraintGraph::removeNode(TypeVariableType *typeVar) {
   // Remove this node.
   auto &impl = typeVar->getImpl();
   unsigned index = impl.getGraphIndex();
-  delete impl.getGraphNode();
+  auto *node = impl.getGraphNode();
+  node->reset();
+  FreeList.push_back(node);
   impl.setGraphNode(nullptr);
 
   // Remove this type variable from the list.