Add more tests for erase(). Fix a few exposed bugs.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@120227 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 68 insertions, 5 deletions

diff --git a/include/llvm/ADT/IntervalMap.h b/include/llvm/ADT/IntervalMap.h
index e1438a3..0451a76 100644
--- a/include/llvm/ADT/IntervalMap.h
+++ b/include/llvm/ADT/IntervalMap.h
@@ -922,6 +922,14 @@ public:
   /// @param Level Move node(Level).
   void moveRight(unsigned Level);
 
+  /// atBegin - Return true if path is at begin().
+  bool atBegin() const {
+    for (unsigned i = 0, e = path.size(); i != e; ++i)
+      if (path[i].offset != 0)
+        return false;
+    return true;
+  }
+
   /// atLastBranch - Return true if the path is at the last branch of the node
   /// at Level.
   /// @param Level Node to examine.
@@ -1523,13 +1531,36 @@ class IntervalMap<KeyT, ValT, N, Traits>::iterator : public const_iterator {
   template <typename NodeT> bool overflow(unsigned Level);
   void treeInsert(KeyT a, KeyT b, ValT y);
   void eraseNode(unsigned Level);
-  void treeErase();
+  void treeErase(bool UpdateRoot = true);
 public:
   /// insert - Insert mapping [a;b] -> y before the current position.
   void insert(KeyT a, KeyT b, ValT y);
 
   /// erase - Erase the current interval.
   void erase();
+
+  iterator &operator++() {
+    const_iterator::operator++();
+    return *this;
+  }
+
+  iterator operator++(int) {
+    iterator tmp = *this;
+    operator++();
+    return tmp;
+  }
+
+  iterator &operator--() {
+    const_iterator::operator--();
+    return *this;
+  }
+
+  iterator operator--(int) {
+    iterator tmp = *this;
+    operator--();
+    return tmp;
+  }
+
 };
 
 /// setNodeStop - Update the stop key of the current node at level and above.
@@ -1660,7 +1691,7 @@ iterator::treeInsert(KeyT a, KeyT b, ValT y) {
           // We have both left and right coalescing. Erase the old SibLeaf entry
           // and continue inserting the larger interval.
           a = SibLeaf.start(SibOfs);
-          erase();
+          treeErase(/* UpdateRoot= */false);
         }
       }
     } else {
@@ -1704,7 +1735,7 @@ iterator::erase() {
 /// treeErase - erase() for a branched tree.
 template <typename KeyT, typename ValT, unsigned N, typename Traits>
 void IntervalMap<KeyT, ValT, N, Traits>::
-iterator::treeErase() {
+iterator::treeErase(bool UpdateRoot) {
   IntervalMap &IM = *this->map;
   IntervalMapImpl::Path &P = this->path;
   Leaf &Node = P.leaf<Leaf>();
@@ -1713,6 +1744,9 @@ iterator::treeErase() {
   if (P.leafSize() == 1) {
     IM.deleteNode(&Node);
     eraseNode(IM.height);
+    // Update rootBranchStart if we erased begin().
+    if (UpdateRoot && IM.branched() && P.valid() && P.atBegin())
+      IM.rootBranchStart() = P.leaf<Leaf>().start(0);
     return;
   }
 
@@ -1724,7 +1758,8 @@ iterator::treeErase() {
   if (P.leafOffset() == NewSize) {
     setNodeStop(IM.height, Node.stop(NewSize - 1));
     P.moveRight(IM.height);
-  }
+  } else if (UpdateRoot && P.atBegin())
+    IM.rootBranchStart() = P.leaf<Leaf>().start(0);
 }
 
 /// eraseNode - Erase the current node at Level from its parent and move path to
diff --git a/unittests/ADT/IntervalMapTest.cpp b/unittests/ADT/IntervalMapTest.cpp
index c065362..4664650 100644
--- a/unittests/ADT/IntervalMapTest.cpp
+++ b/unittests/ADT/IntervalMapTest.cpp
@@ -90,6 +90,10 @@ TEST(IntervalMapTest, SingleEntryMap) {
   EXPECT_EQ(1u, I.value());
   EXPECT_TRUE(I == map.begin());
   EXPECT_FALSE(I == map.end());
+
+  I.erase();
+  EXPECT_TRUE(map.empty());
+  EXPECT_EQ(0, std::distance(map.begin(), map.end()));
 }
 
 // Flat coalescing tests.
@@ -160,6 +164,18 @@ TEST(IntervalMapTest, RootCoalescing) {
   EXPECT_EQ(210u, map.stop());
   EXPECT_EQ(2u, map.lookup(201));
   EXPECT_EQ(1u, map.lookup(200));
+
+  // Erase from the left.
+  map.begin().erase();
+  EXPECT_EQ(2, std::distance(map.begin(), map.end()));
+  EXPECT_EQ(70u, map.start());
+  EXPECT_EQ(210u, map.stop());
+
+  // Erase from the right.
+  (--map.end()).erase();
+  EXPECT_EQ(1, std::distance(map.begin(), map.end()));
+  EXPECT_EQ(70u, map.start());
+  EXPECT_EQ(200u, map.stop());
 }
 
 // Flat multi-coalescing tests.
@@ -330,8 +346,11 @@ TEST(IntervalMapTest, Branched) {
 
   // Insert enough intervals to force a branched tree.
   // This creates 9 leaf nodes with 11 elements each, tree height = 1.
-  for (unsigned i = 1; i < 100; ++i)
+  for (unsigned i = 1; i < 100; ++i) {
     map.insert(10*i, 10*i+5, i);
+    EXPECT_EQ(10u, map.start());
+    EXPECT_EQ(10*i+5, map.stop());
+  }
 
   // Tree limits.
   EXPECT_FALSE(map.empty());
@@ -368,6 +387,15 @@ TEST(IntervalMapTest, Branched) {
   }
   EXPECT_TRUE(I == map.begin());
 
+  // Erase from the front.
+  for (unsigned i = 0; i != 20; ++i) {
+    I.erase();
+    EXPECT_TRUE(I == map.begin());
+    EXPECT_FALSE(map.empty());
+    EXPECT_EQ(I.start(), map.start());
+    EXPECT_EQ(995u, map.stop());
+  }
+
   // Test clear() on branched map.
   map.clear();
   EXPECT_TRUE(map.empty());