Add IntervalMap::iterator::set{Start,Stop,Value} methods that allow limited

editing of the current interval.

These methods may cause coalescing, there are corresponding set*Unchecked
methods for editing without coalescing. The non-coalescing methods are useful
for applying monotonic transforms to all keys or values in a map without
accidentally coalescing transformed and untransformed intervals.

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

3 files changed, 294 insertions, 31 deletions

diff --git a/include/llvm/ADT/IntervalMap.h b/include/llvm/ADT/IntervalMap.h
index b0a0a0ba9a..33569269c6 100644
--- a/include/llvm/ADT/IntervalMap.h
+++ b/include/llvm/ADT/IntervalMap.h
@@ -904,10 +904,10 @@ public:
     return true;
   }
 
-  /// atLastBranch - Return true if the path is at the last branch of the node
-  /// at Level.
+  /// atLastEntry - Return true if the path is at the last entry of the node at
+  /// Level.
   /// @param Level Node to examine.
-  bool atLastBranch(unsigned Level) const {
+  bool atLastEntry(unsigned Level) const {
     return path[Level].offset == path[Level].size - 1;
   }
 
@@ -1365,37 +1365,44 @@ protected:
   void treeFind(KeyT x);
   void treeAdvanceTo(KeyT x);
 
-public:
-  /// const_iterator - Create an iterator that isn't pointing anywhere.
-  const_iterator() : map(0) {}
-
-  /// valid - Return true if the current position is valid, false for end().
-  bool valid() const { return path.valid(); }
-
-  /// start - Return the beginning of the current interval.
-  const KeyT &start() const {
+  /// unsafeStart - Writable access to start() for iterator.
+  KeyT &unsafeStart() const {
     assert(valid() && "Cannot access invalid iterator");
     return branched() ? path.leaf<Leaf>().start(path.leafOffset()) :
                         path.leaf<RootLeaf>().start(path.leafOffset());
   }
 
-  /// stop - Return the end of the current interval.
-  const KeyT &stop() const {
+  /// unsafeStop - Writable access to stop() for iterator.
+  KeyT &unsafeStop() const {
     assert(valid() && "Cannot access invalid iterator");
     return branched() ? path.leaf<Leaf>().stop(path.leafOffset()) :
                         path.leaf<RootLeaf>().stop(path.leafOffset());
   }
 
-  /// value - Return the mapped value at the current interval.
-  const ValT &value() const {
+  /// unsafeValue - Writable access to value() for iterator.
+  ValT &unsafeValue() const {
     assert(valid() && "Cannot access invalid iterator");
     return branched() ? path.leaf<Leaf>().value(path.leafOffset()) :
                         path.leaf<RootLeaf>().value(path.leafOffset());
   }
 
-  const ValT &operator*() const {
-    return value();
-  }
+public:
+  /// const_iterator - Create an iterator that isn't pointing anywhere.
+  const_iterator() : map(0) {}
+
+  /// valid - Return true if the current position is valid, false for end().
+  bool valid() const { return path.valid(); }
+
+  /// start - Return the beginning of the current interval.
+  const KeyT &start() const { return unsafeStart(); }
+
+  /// stop - Return the end of the current interval.
+  const KeyT &stop() const { return unsafeStop(); }
+
+  /// value - Return the mapped value at the current interval.
+  const ValT &value() const { return unsafeValue(); }
+
+  const ValT &operator*() const { return value(); }
 
   bool operator==(const const_iterator &RHS) const {
     assert(map == RHS.map && "Cannot compare iterators from different maps");
@@ -1554,10 +1561,50 @@ class IntervalMap<KeyT, ValT, N, Traits>::iterator : public const_iterator {
   void treeInsert(KeyT a, KeyT b, ValT y);
   void eraseNode(unsigned Level);
   void treeErase(bool UpdateRoot = true);
+  bool canCoalesceLeft(KeyT Start, ValT x);
+  bool canCoalesceRight(KeyT Stop, ValT x);
+
 public:
   /// iterator - Create null iterator.
   iterator() {}
 
+  /// setStart - Move the start of the current interval.
+  /// This may cause coalescing with the previous interval.
+  /// @param a New start key, must not overlap the previous interval.
+  void setStart(KeyT a);
+
+  /// setStop - Move the end of the current interval.
+  /// This may cause coalescing with the following interval.
+  /// @param b New stop key, must not overlap the following interval.
+  void setStop(KeyT b);
+
+  /// setValue - Change the mapped value of the current interval.
+  /// This may cause coalescing with the previous and following intervals.
+  /// @param x New value.
+  void setValue(ValT x);
+
+  /// setStartUnchecked - Move the start of the current interval without
+  /// checking for coalescing or overlaps.
+  /// This should only be used when it is known that coalescing is not required.
+  /// @param a New start key.
+  void setStartUnchecked(KeyT a) { this->unsafeStart() = a; }
+
+  /// setStopUnchecked - Move the end of the current interval without checking
+  /// for coalescing or overlaps.
+  /// This should only be used when it is known that coalescing is not required.
+  /// @param b New stop key.
+  void setStopUnchecked(KeyT b) {
+    this->unsafeStop() = b;
+    // Update keys in branch nodes as well.
+    if (this->path.atLastEntry(this->path.height()))
+      setNodeStop(this->path.height(), b);
+  }
+
+  /// setValueUnchecked - Change the mapped value of the current interval
+  /// without checking for coalescing.
+  /// @param x New value.
+  void setValueUnchecked(ValT x) { this->unsafeValue() = x; }
+
   /// insert - Insert mapping [a;b] -> y before the current position.
   void insert(KeyT a, KeyT b, ValT y);
 
@@ -1588,6 +1635,62 @@ public:
 
 };
 
+/// canCoalesceLeft - Can the current interval coalesce to the left after
+/// changing start or value?
+/// @param Start New start of current interval.
+/// @param Value New value for current interval.
+/// @return True when updating the current interval would enable coalescing.
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+bool IntervalMap<KeyT, ValT, N, Traits>::
+iterator::canCoalesceLeft(KeyT Start, ValT Value) {
+  using namespace IntervalMapImpl;
+  Path &P = this->path;
+  if (!this->branched()) {
+    unsigned i = P.leafOffset();
+    RootLeaf &Node = P.leaf<RootLeaf>();
+    return i && Node.value(i-1) == Value &&
+                Traits::adjacent(Node.stop(i-1), Start);
+  }
+  // Branched.
+  if (unsigned i = P.leafOffset()) {
+    Leaf &Node = P.leaf<Leaf>();
+    return Node.value(i-1) == Value && Traits::adjacent(Node.stop(i-1), Start);
+  } else if (NodeRef NR = P.getLeftSibling(P.height())) {
+    unsigned i = NR.size() - 1;
+    Leaf &Node = NR.get<Leaf>();
+    return Node.value(i) == Value && Traits::adjacent(Node.stop(i), Start);
+  }
+  return false;
+}
+
+/// canCoalesceRight - Can the current interval coalesce to the right after
+/// changing stop or value?
+/// @param Stop New stop of current interval.
+/// @param Value New value for current interval.
+/// @return True when updating the current interval would enable coalescing.
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+bool IntervalMap<KeyT, ValT, N, Traits>::
+iterator::canCoalesceRight(KeyT Stop, ValT Value) {
+  using namespace IntervalMapImpl;
+  Path &P = this->path;
+  unsigned i = P.leafOffset() + 1;
+  if (!this->branched()) {
+    if (i >= P.leafSize())
+      return false;
+    RootLeaf &Node = P.leaf<RootLeaf>();
+    return Node.value(i) == Value && Traits::adjacent(Stop, Node.start(i));
+  }
+  // Branched.
+  if (i < P.leafSize()) {
+    Leaf &Node = P.leaf<Leaf>();
+    return Node.value(i) == Value && Traits::adjacent(Stop, Node.start(i));
+  } else if (NodeRef NR = P.getRightSibling(P.height())) {
+    Leaf &Node = NR.get<Leaf>();
+    return Node.value(0) == Value && Traits::adjacent(Stop, Node.start(0));
+  }
+  return false;
+}
+
 /// setNodeStop - Update the stop key of the current node at level and above.
 template <typename KeyT, typename ValT, unsigned N, typename Traits>
 void IntervalMap<KeyT, ValT, N, Traits>::
@@ -1599,13 +1702,61 @@ iterator::setNodeStop(unsigned Level, KeyT Stop) {
   // Update nodes pointing to the current node.
   while (--Level) {
     P.node<Branch>(Level).stop(P.offset(Level)) = Stop;
-    if (!P.atLastBranch(Level))
+    if (!P.atLastEntry(Level))
       return;
   }
   // Update root separately since it has a different layout.
   P.node<RootBranch>(Level).stop(P.offset(Level)) = Stop;
 }
 
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+void IntervalMap<KeyT, ValT, N, Traits>::
+iterator::setStart(KeyT a) {
+  assert(Traits::stopLess(a, this->stop()) && "Cannot move start beyond stop");
+  KeyT &CurStart = this->unsafeStart();
+  if (!Traits::startLess(a, CurStart) || !canCoalesceLeft(a, this->value())) {
+    CurStart = a;
+    return;
+  }
+  // Coalesce with the interval to the left.
+  --*this;
+  a = this->start();
+  erase();
+  setStartUnchecked(a);
+}
+
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+void IntervalMap<KeyT, ValT, N, Traits>::
+iterator::setStop(KeyT b) {
+  assert(Traits::stopLess(this->start(), b) && "Cannot move stop beyond start");
+  if (Traits::startLess(b, this->stop()) ||
+      !canCoalesceRight(b, this->value())) {
+    setStopUnchecked(b);
+    return;
+  }
+  // Coalesce with interval to the right.
+  KeyT a = this->start();
+  erase();
+  setStartUnchecked(a);
+}
+
+template <typename KeyT, typename ValT, unsigned N, typename Traits>
+void IntervalMap<KeyT, ValT, N, Traits>::
+iterator::setValue(ValT x) {
+  setValueUnchecked(x);
+  if (canCoalesceRight(this->stop(), x)) {
+    KeyT a = this->start();
+    erase();
+    setStartUnchecked(a);
+  }
+  if (canCoalesceLeft(this->start(), x)) {
+    --*this;
+    KeyT a = this->start();
+    erase();
+    setStartUnchecked(a);
+  }
+}
+
 /// insertNode - insert a node before the current path at level.
 /// Leave the current path pointing at the new node.
 /// @param Level path index of the node to be inserted.
@@ -1650,7 +1801,7 @@ iterator::insertNode(unsigned Level, IntervalMapImpl::NodeRef Node, KeyT Stop) {
   }
   P.node<Branch>(Level).insert(P.offset(Level), P.size(Level), Node, Stop);
   P.setSize(Level, P.size(Level) + 1);
-  if (P.atLastBranch(Level))
+  if (P.atLastEntry(Level))
     setNodeStop(Level, Stop);
   P.reset(Level + 1);
   return SplitRoot;
diff --git a/lib/Support/IntervalMap.cpp b/lib/Support/IntervalMap.cpp
index 6f39b183a1..4dfcc404ca 100644
--- a/lib/Support/IntervalMap.cpp
+++ b/lib/Support/IntervalMap.cpp
@@ -79,11 +79,11 @@ NodeRef Path::getRightSibling(unsigned Level) const {
 
   // Go up the tree until we can go right.
   unsigned l = Level - 1;
-  while (l && atLastBranch(l))
+  while (l && atLastEntry(l))
     --l;
 
   // We can't go right.
-  if (atLastBranch(l))
+  if (atLastEntry(l))
     return NodeRef();
 
   // NR is the subtree containing our right sibling.
@@ -100,7 +100,7 @@ void Path::moveRight(unsigned Level) {
 
   // Go up the tree until we can go right.
   unsigned l = Level - 1;
-  while (l && atLastBranch(l))
+  while (l && atLastEntry(l))
     --l;
 
   // NR is the subtree containing our right sibling. If we hit end(), we have
diff --git a/unittests/ADT/IntervalMapTest.cpp b/unittests/ADT/IntervalMapTest.cpp
index 445afcaf7a..fc16a32798 100644
--- a/unittests/ADT/IntervalMapTest.cpp
+++ b/unittests/ADT/IntervalMapTest.cpp
@@ -14,8 +14,7 @@ using namespace llvm;
 
 namespace {
 
-typedef IntervalMap<unsigned, unsigned> UUMap;
-typedef IntervalMap<unsigned, unsigned, 4> UU4Map;
+typedef IntervalMap<unsigned, unsigned, 4> UUMap;
 
 // Empty map tests
 TEST(IntervalMapTest, EmptyMap) {
@@ -99,6 +98,40 @@ TEST(IntervalMapTest, SingleEntryMap) {
   EXPECT_TRUE(I == map.begin());
   EXPECT_FALSE(I == map.end());
 
+  // Change the value.
+  I.setValue(2);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(100u, I.start());
+  EXPECT_EQ(150u, I.stop());
+  EXPECT_EQ(2u, I.value());
+
+  // Grow the bounds.
+  I.setStart(0);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(0u, I.start());
+  EXPECT_EQ(150u, I.stop());
+  EXPECT_EQ(2u, I.value());
+
+  I.setStop(200);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(0u, I.start());
+  EXPECT_EQ(200u, I.stop());
+  EXPECT_EQ(2u, I.value());
+
+  // Shrink the bounds.
+  I.setStart(150);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(150u, I.start());
+  EXPECT_EQ(200u, I.stop());
+  EXPECT_EQ(2u, I.value());
+
+  I.setStop(160);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(150u, I.start());
+  EXPECT_EQ(160u, I.stop());
+  EXPECT_EQ(2u, I.value());
+
+  // Erase last elem.
   I.erase();
   EXPECT_TRUE(map.empty());
   EXPECT_EQ(0, std::distance(map.begin(), map.end()));
@@ -160,6 +193,18 @@ TEST(IntervalMapTest, RootCoalescing) {
   EXPECT_EQ(1, std::distance(map.begin(), map.end()));
   EXPECT_EQ(90u, map.start());
   EXPECT_EQ(200u, map.stop());
+
+  // Add non-coalescing, then trigger coalescing with setValue.
+  map.insert(80, 89, 2);
+  map.insert(201, 210, 2);
+  EXPECT_EQ(3, std::distance(map.begin(), map.end()));
+  (++map.begin()).setValue(2);
+  EXPECT_EQ(1, std::distance(map.begin(), map.end()));
+  I = map.begin();
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(80u, I.start());
+  EXPECT_EQ(210u, I.stop());
+  EXPECT_EQ(2u, I.value());
 }
 
 // Flat multi-coalescing tests.
@@ -324,12 +369,79 @@ TEST(IntervalMapTest, Branched) {
   EXPECT_EQ(20u, I.start());
   EXPECT_EQ(25u, I.stop());
 
+  // Change value, no coalescing.
+  I.setValue(0);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(20u, I.start());
+  EXPECT_EQ(25u, I.stop());
+  EXPECT_EQ(0u, I.value());
+
+  // Close the gap right, no coalescing.
+  I.setStop(29);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(20u, I.start());
+  EXPECT_EQ(29u, I.stop());
+  EXPECT_EQ(0u, I.value());
+
+  // Change value, no coalescing.
+  I.setValue(2);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(20u, I.start());
+  EXPECT_EQ(29u, I.stop());
+  EXPECT_EQ(2u, I.value());
+
+  // Change value, now coalescing.
+  I.setValue(3);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(20u, I.start());
+  EXPECT_EQ(35u, I.stop());
+  EXPECT_EQ(3u, I.value());
+
+  // Close the gap, now coalescing.
+  I.setValue(4);
+  ASSERT_TRUE(I.valid());
+  I.setStop(39);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(20u, I.start());
+  EXPECT_EQ(45u, I.stop());
+  EXPECT_EQ(4u, I.value());
+
   // advanceTo another node.
   I.advanceTo(200);
   ASSERT_TRUE(I.valid());
   EXPECT_EQ(200u, I.start());
   EXPECT_EQ(205u, I.stop());
 
+  // Close the gap left, no coalescing.
+  I.setStart(196);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(196u, I.start());
+  EXPECT_EQ(205u, I.stop());
+  EXPECT_EQ(20u, I.value());
+
+  // Change value, no coalescing.
+  I.setValue(0);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(196u, I.start());
+  EXPECT_EQ(205u, I.stop());
+  EXPECT_EQ(0u, I.value());
+
+  // Change value, now coalescing.
+  I.setValue(19);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(190u, I.start());
+  EXPECT_EQ(205u, I.stop());
+  EXPECT_EQ(19u, I.value());
+
+  // Close the gap, now coalescing.
+  I.setValue(18);
+  ASSERT_TRUE(I.valid());
+  I.setStart(186);
+  ASSERT_TRUE(I.valid());
+  EXPECT_EQ(180u, I.start());
+  EXPECT_EQ(205u, I.stop());
+  EXPECT_EQ(18u, I.value());
+
   // Erase from the front.
   I = map.begin();
   for (unsigned i = 0; i != 20; ++i) {
@@ -348,8 +460,8 @@ TEST(IntervalMapTest, Branched) {
 
 // Branched, high, non-coalescing tests.
 TEST(IntervalMapTest, Branched2) {
-  UU4Map::Allocator allocator;
-  UU4Map map(allocator);
+  UUMap::Allocator allocator;
+  UUMap map(allocator);
 
   // Insert enough intervals to force a height >= 2 tree.
   for (unsigned i = 1; i < 1000; ++i)
@@ -369,7 +481,7 @@ TEST(IntervalMapTest, Branched2) {
   }
 
   // Forward iteration.
-  UU4Map::iterator I = map.begin();
+  UUMap::iterator I = map.begin();
   for (unsigned i = 1; i < 1000; ++i) {
     ASSERT_TRUE(I.valid());
     EXPECT_EQ(10*i, I.start());
@@ -416,8 +528,8 @@ TEST(IntervalMapTest, Branched2) {
 
 // Random insertions, coalescing to a single interval.
 TEST(IntervalMapTest, RandomCoalescing) {
-  UU4Map::Allocator allocator;
-  UU4Map map(allocator);
+  UUMap::Allocator allocator;
+  UUMap map(allocator);
 
   // This is a poor PRNG with maximal period:
   // x_n = 5 x_{n-1} + 1 mod 2^N