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.

3 files changed, 33 insertions, 10 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++;
diff --git a/unittests/Searcher/SearcherTest.cpp b/unittests/Searcher/SearcherTest.cpp
index fdd93e62..eff5a8af 100644
--- a/unittests/Searcher/SearcherTest.cpp
+++ b/unittests/Searcher/SearcherTest.cpp
@@ -51,13 +51,13 @@ TEST(SearcherTest, RandomPath) {
   }
 
   rp.update(&es, {&es1}, {&es});
-  PTree::remove(es.ptreeNode);
+  processTree.remove(es.ptreeNode);
   for (int i = 0; i < 100; i++) {
     EXPECT_EQ(&rp.selectState(), &es1);
   }
 
   rp.update(&es1, {}, {&es1});
-  PTree::remove(es1.ptreeNode);
+  processTree.remove(es1.ptreeNode);
   EXPECT_TRUE(rp.empty());
 }
 
@@ -103,16 +103,16 @@ TEST(SearcherTest, TwoRandomPath) {
   }
 
   rp1.update(&es, {}, {&es});
-  PTree::remove(es.ptreeNode);
+  processTree.remove(es.ptreeNode);
   EXPECT_TRUE(rp1.empty());
   EXPECT_EQ(&rp.selectState(), &es1);
 
   rp.update(&es1, {}, {&es1});
-  PTree::remove(es1.ptreeNode);
+  processTree.remove(es1.ptreeNode);
   EXPECT_TRUE(rp.empty());
   EXPECT_TRUE(rp1.empty());
 
-  PTree::remove(root.ptreeNode);
+  processTree.remove(root.ptreeNode);
 }
 
 TEST(SearcherTest, TwoRandomPathDot) {
@@ -174,7 +174,7 @@ TEST(SearcherTest, TwoRandomPathDot) {
   rp1.update(nullptr, {&es}, {&es1});
 
   rp1.update(&es, {}, {&es});
-  PTree::remove(es.ptreeNode);
+  processTree.remove(es.ptreeNode);
 
   modelPTreeDot.str("");
   modelPTreeDot
@@ -198,15 +198,15 @@ TEST(SearcherTest, TwoRandomPathDot) {
   pTreeDot = "";
   processTree.dump(pTreeDotStream);
   EXPECT_EQ(modelPTreeDot.str(), pTreeDotStream.str());
-  PTree::remove(es1.ptreeNode);
-  PTree::remove(root.ptreeNode);
+  processTree.remove(es1.ptreeNode);
+  processTree.remove(root.ptreeNode);
 }
 TEST(SearcherDeathTest, TooManyRandomPaths) {
   // First state
   ExecutionState es;
   PTree processTree(&es);
   es.ptreeNode = processTree.root.getPointer();
-  PTree::remove(es.ptreeNode); // Need to remove to avoid leaks
+  processTree.remove(es.ptreeNode); // Need to remove to avoid leaks
 
   RNG rng;
   RandomPathSearcher rp(processTree, rng);