Implement IntervalMap::clear().

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

2 files changed, 34 insertions, 3 deletions

diff --git a/include/llvm/ADT/IntervalMap.h b/include/llvm/ADT/IntervalMap.h
index dd4eec3c63..96b5161da7 100644
--- a/include/llvm/ADT/IntervalMap.h
+++ b/include/llvm/ADT/IntervalMap.h
@@ -806,7 +806,7 @@ private:
   Leaf *allocLeaf()  {
     return new(allocator.template Allocate<Leaf>()) Leaf();
   }
-  void freeLeaf(Leaf *P) {
+  void deleteLeaf(Leaf *P) {
     P->~Leaf();
     allocator.Deallocate(P);
   }
@@ -814,7 +814,7 @@ private:
   Branch *allocBranch() {
     return new(allocator.template Allocate<Branch>()) Branch();
   }
-  void freeBranch(Branch *P) {
+  void deleteBranch(Branch *P) {
     P->~Branch();
     allocator.Deallocate(P);
   }
@@ -838,8 +838,8 @@ private:
   bool branched() const { return height > 0; }
 
   ValT treeSafeLookup(KeyT x, ValT NotFound) const;
-
   void visitNodes(void (IntervalMap::*f)(NodeRef, unsigned Level));
+  void deleteNode(NodeRef Node, unsigned Level);
 
 public:
   explicit IntervalMap(Allocator &a) : height(0), rootSize(0), allocator(a) {
@@ -881,6 +881,9 @@ public:
     find(a).insert(a, b, y);
   }
 
+  /// clear - Remove all entries.
+  void clear();
+
   class const_iterator;
   class iterator;
   friend class const_iterator;
@@ -1048,6 +1051,25 @@ visitNodes(void (IntervalMap::*f)(NodeRef, unsigned Height)) {
     (this->*f)(Refs[i], 0);
 }
 
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+void IntervalMap<KeyT, ValT, N, Traits>::
+deleteNode(NodeRef Node, unsigned Level) {
+  if (Level)
+    deleteBranch(&Node.branch());
+  else
+    deleteLeaf(&Node.leaf());
+}
+
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+void IntervalMap<KeyT, ValT, N, Traits>::
+clear() {
+  if (branched()) {
+    visitNodes(&IntervalMap::deleteNode);
+    switchRootToLeaf();
+  }
+  rootSize = 0;
+}
+
 #ifndef NDEBUG
 template <typename KeyT, typename ValT, unsigned N, typename Traits>
 void IntervalMap<KeyT, ValT, N, Traits>::
diff --git a/unittests/ADT/IntervalMapTest.cpp b/unittests/ADT/IntervalMapTest.cpp
index c7def84340..d4b2f52b50 100644
--- a/unittests/ADT/IntervalMapTest.cpp
+++ b/unittests/ADT/IntervalMapTest.cpp
@@ -315,6 +315,11 @@ TEST(IntervalMapTest, RootMultiCoalescing) {
   EXPECT_EQ(320u, I.stop());
   ++I;
   EXPECT_FALSE(I.valid());
+
+  // Test clear() on non-branched map.
+  map.clear();
+  EXPECT_TRUE(map.empty());
+  EXPECT_TRUE(map.begin() == map.end());
 }
 
 // Branched, non-coalescing tests.
@@ -362,6 +367,10 @@ TEST(IntervalMapTest, Branched) {
   }
   EXPECT_TRUE(I == map.begin());
 
+  // Test clear() on branched map.
+  map.clear();
+  EXPECT_TRUE(map.empty());
+  EXPECT_TRUE(map.begin() == map.end());
 }
 
 } // namespace