diff --git a/include/llvm/ADT/IntervalMap.h b/include/llvm/ADT/IntervalMap.h index 2758cbc562d..bf8dee6f3ae 100644 --- a/include/llvm/ADT/IntervalMap.h +++ b/include/llvm/ADT/IntervalMap.h @@ -869,6 +869,11 @@ public: path.push_back(Entry(Node, Offset)); } + /// pop - Remove the last path entry. + void pop() { + path.pop_back(); + } + /// setSize - Set the size of a node both in the path and in the tree. /// @param Level 0..height. Note that setting the root size won't change /// map->rootSize. @@ -1389,6 +1394,7 @@ protected: void pathFillFind(KeyT x); void treeFind(KeyT x); + void treeAdvanceTo(KeyT x); public: /// valid - Return true if the current position is valid, false for end(). @@ -1497,7 +1503,8 @@ public: }; -// pathFillFind - Complete path by searching for x. +/// pathFillFind - Complete path by searching for x. +/// @param x Key to search for. template void IntervalMap:: const_iterator::pathFillFind(KeyT x) { @@ -1510,7 +1517,8 @@ const_iterator::pathFillFind(KeyT x) { path.push(NR, NR.get().safeFind(0, x)); } -// treeFind - Find in a branched tree. +/// treeFind - Find in a branched tree. +/// @param x Key to search for. template void IntervalMap:: const_iterator::treeFind(KeyT x) { @@ -1519,6 +1527,43 @@ const_iterator::treeFind(KeyT x) { pathFillFind(x); } +/// treeAdvanceTo - Find position after the current one. +/// @param x Key to search for. +template +void IntervalMap:: +const_iterator::treeAdvanceTo(KeyT x) { + // Can we stay on the same leaf node? + if (!Traits::stopLess(path.leaf().stop(path.leafSize() - 1), x)) { + path.leafOffset() = path.leaf().safeFind(path.leafOffset(), x); + return; + } + + // Drop the current leaf. + path.pop(); + + // Search towards the root for a usable subtree. + if (path.height()) { + for (unsigned l = path.height() - 1; l; --l) { + if (!Traits::stopLess(path.node(l).stop(path.offset(l)), x)) { + // The branch node at l+1 is usable + path.offset(l + 1) = + path.node(l + 1).safeFind(path.offset(l + 1), x); + return pathFillFind(x); + } + path.pop(); + } + // Is the level-1 Branch usable? + if (!Traits::stopLess(map->rootBranch().stop(path.offset(0)), x)) { + path.offset(1) = path.node(1).safeFind(path.offset(1), x); + return pathFillFind(x); + } + } + + // We reached the root. + setRoot(map->rootBranch().findFrom(path.offset(0), map->rootSize, x)); + if (valid()) + pathFillFind(x); +} //===----------------------------------------------------------------------===// //--- iterator ----// diff --git a/unittests/ADT/IntervalMapTest.cpp b/unittests/ADT/IntervalMapTest.cpp index 466465064ab..f4b1ebc8d3a 100644 --- a/unittests/ADT/IntervalMapTest.cpp +++ b/unittests/ADT/IntervalMapTest.cpp @@ -206,6 +206,17 @@ TEST(IntervalMapTest, RootMultiCoalescing) { ++I; EXPECT_FALSE(I.valid()); + // Test advanceTo on flat tree. + I = map.begin(); + I.advanceTo(135); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(140u, I.start()); + EXPECT_EQ(150u, I.stop()); + + I.advanceTo(145); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(140u, I.start()); + EXPECT_EQ(150u, I.stop()); // Coalesce left with followers. // [100;110] [120;130] [140;150] [160;170] @@ -387,7 +398,20 @@ TEST(IntervalMapTest, Branched) { } EXPECT_TRUE(I == map.begin()); + // Test advanceTo in same node. + I.advanceTo(20); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(20u, I.start()); + EXPECT_EQ(25u, I.stop()); + + // advanceTo another node. + I.advanceTo(200); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(200u, I.start()); + EXPECT_EQ(205u, I.stop()); + // Erase from the front. + I = map.begin(); for (unsigned i = 0; i != 20; ++i) { I.erase(); EXPECT_TRUE(I == map.begin()); @@ -446,6 +470,24 @@ TEST(IntervalMapTest, Branched2) { } EXPECT_TRUE(I == map.begin()); + // Test advanceTo in same node. + I.advanceTo(20); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(20u, I.start()); + EXPECT_EQ(25u, I.stop()); + + // advanceTo sibling leaf node. + I.advanceTo(200); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(200u, I.start()); + EXPECT_EQ(205u, I.stop()); + + // advanceTo further. + I.advanceTo(2000); + ASSERT_TRUE(I.valid()); + EXPECT_EQ(2000u, I.start()); + EXPECT_EQ(2005u, I.stop()); + // Test clear() on branched map. map.clear(); EXPECT_TRUE(map.empty());