Fix PTree::remove to clean the tree properly

The previous version left unnecessary intermediate nodes behind, sometimes
leading to very long paths in the tree.

2 files changed, 24 insertions, 1 deletions

diff --git a/lib/Core/PTree.cpp b/lib/Core/PTree.cpp
index 244e8520..4d2931cb 100644
--- a/lib/Core/PTree.cpp
+++ b/lib/Core/PTree.cpp
@@ -54,6 +54,29 @@ void PTree::remove(PTreeNode *n) {
     delete n;
     n = p;
   } while (n && !n->left.getPointer() && !n->right.getPointer());
+
+  if (n) {
+    // We're now at a node that has exactly one child; we've just deleted the
+    // other one. Eliminate the node and connect its child to the parent
+    // directly (if it's not the root).
+    PTreeNodePtr child = n->left.getPointer() ? n->left : n->right;
+    PTreeNode *parent = n->parent;
+
+    child.getPointer()->parent = parent;
+    if (!parent) {
+      // We're at the root.
+      root = child;
+    } else {
+      if (n == parent->left.getPointer()) {
+        parent->left = child;
+      } else {
+        assert(n == parent->right.getPointer());
+        parent->right = child;
+      }
+    }
+
+    delete n;
+  }
 }
 
 void PTree::dump(llvm::raw_ostream &os) {
diff --git a/lib/Core/PTree.h b/lib/Core/PTree.h
index b1b6e9fc..3d87611d 100644
--- a/lib/Core/PTree.h
+++ b/lib/Core/PTree.h
@@ -49,7 +49,7 @@ namespace klee {
 
     void attach(PTreeNode *node, ExecutionState *leftState,
                 ExecutionState *rightState);
-    static void remove(PTreeNode *node);
+    void remove(PTreeNode *node);
     void dump(llvm::raw_ostream &os);
     std::uint8_t getNextId() {
       std::uint8_t id = 1 << registeredIds++;