Fixed erase return iterator.

5cript · Jul 31, 2024 · 3936ada · 3936ada
1 parent 0f2830a
commit 3936ada
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 11 deletions.
diff --git a/include/interval-tree/interval_tree.hpp b/include/interval-tree/interval_tree.hpp
@@ -886,19 +886,27 @@ namespace lib_interval_tree
                 throw std::out_of_range("cannot erase end iterator");
 
             auto next = iter;
-            ++next;
 
-            node_type* y;
-            if (!iter.node_->left_ || !iter.node_->right_)
-                y = iter.node_;
-            else
-                y = successor(iter.node_);
+            node_type* y = [&next, &iter, this]() {
+                if (!iter.node_->left_ || !iter.node_->right_)
+                {
+                    ++next;
+                    return iter.node_;
+                }
+                else
+                {
+                    const auto y = successor(iter.node_);
+                    next = iterator{iter.node_, this};
+                    return y;
+                }
+            }();
 
-            node_type* x;
-            if (y->left_)
-                x = y->left_;
-            else
-                x = y->right_;
+            node_type* x = [y](){
+                if (y->left_)
+                    return y->left_;
+                else
+                    return y->right_;
+            }();
 
             if (x)
                 x->parent_ = y->parent_;

diff --git a/tests/erase_tests.hpp b/tests/erase_tests.hpp
@@ -184,6 +184,25 @@ TEST_F(EraseTests, RandomEraseTest)
     testTreeHeightHealth(tree);
 }
 
+
+
+TEST_F(EraseTests, MassiveDeleteEntireTreeWithEraseReturnIterator)
+{
+    constexpr int amount = 1000;
+
+    for (int i = 0; i != amount; ++i)
+        tree.insert(makeSafeOracleInterval(&oracle, distSmall(gen), distSmall(gen)));
+
+    for(auto iter = tree.begin(); !tree.empty();)
+    {
+        iter = tree.erase(iter);
+    }
+
+    EXPECT_EQ(oracle.livingInstances, 0);
+    testMaxProperty(tree);
+    testTreeHeightHealth(tree);
+}
+
 TEST_F(EraseTests, ReturnedIteratorPointsToNextInOrderNode)
 {
     auto regularTree = makeTree();
@@ -234,3 +253,16 @@ TEST_F(EraseTests, CanEraseEntireTreeUsingReturnedIterator)
         iter = tree.erase(iter);
     EXPECT_EQ(tree.empty(), true);
 }
+
+TEST_F(EraseTests, FromNuiTest)
+{
+    lib_interval_tree::interval_tree_t <int> tree;
+    tree.insert({0, 0});
+    tree.insert({4, 4});
+    tree.insert({13, 13});
+
+    drawTree("erase_tests_from_nui", tree);
+
+    auto iter = tree.erase(tree.find({4, 4}));
+    EXPECT_EQ(*iter, (decltype(tree)::interval_type{13, 13})) << *iter;
+}