[frozen-multimap] Add support for erasing once we have finished constructing the map.

include/swift/Basic/FrozenMultiMap.h

@@ -12,14 +12,24 @@
 ///
 /// \file
 ///
-/// A 2 stage multi-map. Initially the multimap is mutable and can only be
-/// initialized. Once complete, the map is frozen and can be only used for map
-/// operations. It is guaranteed that all values are still in insertion order.
+/// A 2 stage multi-map for use in contexts where one is iterating over an IR in
+/// a read only way and then wants to use a multi-map for post-processing that
+/// can support iteration in insertion order.
 ///
-/// DISCUSSION: These restrictions flow from the internal implementation of the
-/// multi-map being a pair of keys, values. We form the map property by
-/// performing a stable_sort of the (key, value) in the process of freezing the
-/// map.
+/// In the first stage, the multi-map can be initialized with new elements in a
+/// purely additive fashion, but methods that rely on the map being frozen
+/// (find, getRange(), erase, etc) assert.
+///
+/// Once the user has finished iterating over the IR, the map is frozen and then
+/// can only be used for map operations (find), erasing operations (erase), and
+/// deterministic range iteration.
+///
+/// This is accomplished by the internal implementation of the frozen multi map
+/// just being an array of (key, value) pairs that when we freeze, we sort using
+/// a stable sort on the key. Since we use a stable sort on the key, we know
+/// that all of the individual multimap value arrays are still in insertion
+/// order helping us avoid non-determinism like one must deal with when using
+/// other maps.
 ///
 //===----------------------------------------------------------------------===//
 
@@ -34,17 +44,19 @@
 namespace swift {
 
 template <typename Key, typename Value,
-          typename VectorStorage = std::vector<std::pair<Key, Value>>>
+          typename VectorStorage = std::vector<std::pair<Key, Optional<Value>>>>
 class FrozenMultiMap {
   VectorStorage storage;
   bool frozen = false;
 
 private:
   struct PairToSecondElt;
+  struct PairWithTypeErasedOptionalSecondElt;
 
 public:
   using PairToSecondEltRange =
-      TransformRange<ArrayRef<std::pair<Key, Value>>, PairToSecondElt>;
+      TransformRange<ArrayRef<std::pair<Key, Optional<Value>>>,
+                     PairToSecondElt>;
 
   FrozenMultiMap() = default;
 
@@ -60,30 +72,61 @@ class FrozenMultiMap {
     // inst as the first element.
     auto start = std::lower_bound(
         storage.begin(), storage.end(), std::make_pair(key, Value()),
-        [&](const std::pair<Key, Value> &p1, const std::pair<Key, Value> &p2) {
+        [&](const std::pair<Key, Optional<Value>> &p1,
+            const std::pair<Key, Optional<Value>> &p2) {
           return p1.first < p2.first;
         });
-    if (start == storage.end() || start->first != key) {
+
+    // If we did not find that first element or that key has a .none value
+    // (signaling that we erased it), return None.
+    if (start == storage.end() || start->first != key ||
+        !start->second.hasValue()) {
       return None;
     }
 
     // Ok, we found our first element. Now scan forward until we find a pair
     // whose instruction is not our own instruction.
-    auto end = find_if_not(
-        start, storage.end(),
-        [&](const std::pair<Key, Value> &pair) { return pair.first == key; });
+    auto end = find_if_not(start, storage.end(),
+                           [&](const std::pair<Key, Optional<Value>> &pair) {
+                             return pair.first == key;
+                           });
     unsigned count = std::distance(start, end);
-    ArrayRef<std::pair<Key, Value>> slice(&*start, count);
+    ArrayRef<std::pair<Key, Optional<Value>>> slice(&*start, count);
     return PairToSecondEltRange(slice, PairToSecondElt());
   }
 
+  bool erase(const Key &key) {
+    assert(isFrozen() &&
+           "Can not perform an erase operation until the map is frozen");
+    // Since our array is sorted, we need to first find the first pair with our
+    // inst as the first element.
+    auto start = std::lower_bound(
+        storage.begin(), storage.end(), std::make_pair(key, Value()),
+        [&](const std::pair<Key, Optional<Value>> &p1,
+            const std::pair<Key, Optional<Value>> &p2) {
+          return p1.first < p2.first;
+        });
+
+    // If we did not find that first element or that key has a .none value
+    // (signaling that we erased it), return false.
+    if (start == storage.end() || start->first != key ||
+        !start->second.hasValue()) {
+      return false;
+    }
+
+    // Ok, we found our element. Just set its value to .none to signal it was
+    // destroyed and then return true.
+    start->second = None;
+    return true;
+  }
+
   bool isFrozen() const { return frozen; }
 
   /// Set this map into its frozen state when we
   void setFrozen() {
     std::stable_sort(storage.begin(), storage.end(),
-                     [&](const std::pair<Key, Value> &lhs,
-                         const std::pair<Key, Value> &rhs) {
+                     [&](const std::pair<Key, Optional<Value>> &lhs,
+                         const std::pair<Key, Optional<Value>> &rhs) {
                        // Only compare the first entry so that we preserve
                        // insertion order.
                        return lhs.first < rhs.first;
@@ -100,13 +143,14 @@ class FrozenMultiMap {
   unsigned size() const { return storage.size(); }
   bool empty() const { return storage.empty(); }
 
-  struct iterator : std::iterator<std::forward_iterator_tag,
-                                  std::pair<Key, PairToSecondEltRange>> {
+  struct iterator
+      : std::iterator<std::forward_iterator_tag,
+                      std::pair<Key, Optional<PairToSecondEltRange>>> {
     using base_iterator = typename decltype(storage)::iterator;
 
     FrozenMultiMap &map;
     base_iterator baseIter;
-    Optional<std::pair<Key, PairToSecondEltRange>> currentValue;
+    Optional<std::pair<Key, Optional<PairToSecondEltRange>>> currentValue;
 
     iterator(FrozenMultiMap &map, base_iterator iter)
         : map(map), baseIter(iter), currentValue() {
@@ -123,21 +167,27 @@ class FrozenMultiMap {
       }
 
       // Otherwise, determine the next range that we are visiting.
-      auto rangeEnd = std::find_if_not(std::next(baseIter), end,
-                                       [&](const std::pair<Key, Value> &elt) {
-                                         return elt.first == baseIter->first;
-                                       });
-      unsigned count = std::distance(baseIter, rangeEnd);
-      ArrayRef<std::pair<Key, Value>> slice(&*baseIter, count);
-      currentValue = {baseIter->first,
-                      PairToSecondEltRange(slice, PairToSecondElt())};
+      auto rangeEnd =
+          std::find_if_not(std::next(baseIter), end,
+                           [&](const std::pair<Key, Optional<Value>> &elt) {
+                             return elt.first == baseIter->first;
+                           });
+
+      Optional<PairToSecondEltRange> resultRange;
+      if (baseIter->second.hasValue()) {
+        unsigned count = std::distance(baseIter, rangeEnd);
+        ArrayRef<std::pair<Key, Optional<Value>>> slice(&*baseIter, count);
+        resultRange.emplace(slice, PairToSecondElt());
+      }
+      currentValue = {baseIter->first, resultRange};
     }
 
     iterator &operator++() {
-      baseIter = std::find_if_not(std::next(baseIter), map.storage.end(),
-                                  [&](const std::pair<Key, Value> &elt) {
-                                    return elt.first == baseIter->first;
-                                  });
+      baseIter =
+          std::find_if_not(std::next(baseIter), map.storage.end(),
+                           [&](const std::pair<Key, Optional<Value>> &elt) {
+                             return elt.first == baseIter->first;
+                           });
       updateCurrentValue();
       return *this;
     }
@@ -152,7 +202,7 @@ class FrozenMultiMap {
       return tmp;
     }
 
-    std::pair<Key, PairToSecondEltRange> operator*() const {
+    std::pair<Key, Optional<PairToSecondEltRange>> operator*() const {
       return *currentValue;
     }
 
@@ -165,7 +215,19 @@ class FrozenMultiMap {
     }
   };
 
-  using RangeType = llvm::iterator_range<iterator>;
+  struct ToNonErasedValues {
+    Optional<std::pair<Key, Optional<PairToSecondEltRange>>>
+    operator()(std::pair<Key, Optional<PairToSecondEltRange>> pair) const {
+      if (!pair.second.hasValue())
+        return None;
+      return pair;
+    }
+  };
+
+  using IgnoringErasedValueRangeType =
+      OptionalTransformRange<llvm::iterator_range<iterator>, ToNonErasedValues>;
+  using RangeType = TransformRange<IgnoringErasedValueRangeType,
+                                   PairWithTypeErasedOptionalSecondElt>;
 
   /// Return a range of (key, ArrayRef<Value>) pairs. The keys are guaranteed to
   /// be in key sorted order and the ArrayRef<Value> are in insertion order.
@@ -175,22 +237,36 @@ class FrozenMultiMap {
     auto *self = const_cast<FrozenMultiMap *>(this);
     iterator iter1 = iterator(*self, self->storage.begin());
     iterator iter2 = iterator(*self, self->storage.end());
-    return llvm::make_range(iter1, iter2);
+    auto baseRange = llvm::make_range(iter1, iter2);
+    auto optRange = makeOptionalTransformRange(baseRange, ToNonErasedValues());
+    return makeTransformRange(optRange, PairWithTypeErasedOptionalSecondElt());
   }
 };
 
 template <typename Key, typename Value, typename Storage>
 struct FrozenMultiMap<Key, Value, Storage>::PairToSecondElt {
   PairToSecondElt() {}
 
-  Value operator()(const std::pair<Key, Value> &pair) const {
-    return pair.second;
+  Value operator()(const std::pair<Key, Optional<Value>> &pair) const {
+    return *pair.second;
+  }
+};
+
+template <typename Key, typename Value, typename Storage>
+struct FrozenMultiMap<Key, Value,
+                      Storage>::PairWithTypeErasedOptionalSecondElt {
+  PairWithTypeErasedOptionalSecondElt() {}
+
+  std::pair<Key, PairToSecondEltRange>
+  operator()(const std::pair<Key, Optional<PairToSecondEltRange>> &pair) const {
+    return std::make_pair(pair.first, *pair.second);
   }
 };
 
 template <typename Key, typename Value, unsigned SmallSize>
 using SmallFrozenMultiMap =
-    FrozenMultiMap<Key, Value, SmallVector<std::pair<Key, Value>, SmallSize>>;
+    FrozenMultiMap<Key, Value,
+                   SmallVector<std::pair<Key, Optional<Value>>, SmallSize>>;
 
 } // namespace swift
 

unittests/Basic/FrozenMultiMapTest.cpp

@@ -223,3 +223,101 @@ TEST(FrozenMultiMapCustomTest, RandomAgainstStdMultiMap) {
     }
   }
 }
+
+TEST(FrozenMultiMapCustomTest, SimpleErase1) {
+  Canary::resetIDs();
+  FrozenMultiMap<Canary, Canary> map;
+
+  auto key1 = Canary();
+  auto key2 = Canary();
+  map.insert(key1, Canary());
+  map.insert(key1, Canary());
+  map.insert(key1, Canary());
+  map.insert(key2, Canary());
+  map.insert(key2, Canary());
+
+  map.setFrozen();
+
+  EXPECT_EQ(map.size(), 5u);
+
+  EXPECT_TRUE(map.erase(key2));
+  EXPECT_FALSE(map.erase(key2));
+
+  {
+    auto range = map.getRange();
+    EXPECT_EQ(std::distance(range.begin(), range.end()), 1);
+
+    {
+      auto begin = range.begin();
+      auto end = range.end();
+      ++begin;
+      EXPECT_EQ(std::distance(begin, end), 0);
+    }
+
+    auto iter = range.begin();
+    {
+      auto p = *iter;
+      EXPECT_EQ(p.first.getID(), key1.getID());
+      EXPECT_EQ(p.second.size(), 3u);
+      EXPECT_EQ(p.second[0].getID(), 2u);
+      EXPECT_EQ(p.second[1].getID(), 3u);
+      EXPECT_EQ(p.second[2].getID(), 4u);
+    }
+  }
+
+  EXPECT_TRUE(map.erase(key1));
+  EXPECT_FALSE(map.erase(key1));
+
+  {
+    auto range = map.getRange();
+    EXPECT_EQ(std::distance(range.begin(), range.end()), 0);
+  }
+}
+
+TEST(FrozenMultiMapCustomTest, SimpleErase2) {
+  Canary::resetIDs();
+  FrozenMultiMap<Canary, Canary> map;
+
+  auto key1 = Canary();
+  auto key2 = Canary();
+  map.insert(key1, Canary());
+  map.insert(key1, Canary());
+  map.insert(key1, Canary());
+  map.insert(key2, Canary());
+  map.insert(key2, Canary());
+
+  map.setFrozen();
+
+  EXPECT_EQ(map.size(), 5u);
+
+  EXPECT_TRUE(map.erase(key1));
+
+  {
+    auto range = map.getRange();
+    EXPECT_EQ(std::distance(range.begin(), range.end()), 1);
+
+    {
+      auto begin = range.begin();
+      auto end = range.end();
+      ++begin;
+      EXPECT_EQ(std::distance(begin, end), 0);
+    }
+
+    auto iter = range.begin();
+    {
+      auto p = *iter;
+      EXPECT_EQ(p.first.getID(), key2.getID());
+      EXPECT_EQ(p.second.size(), 2u);
+      EXPECT_EQ(p.second[0].getID(), 5u);
+      EXPECT_EQ(p.second[1].getID(), 6u);
+    }
+  }
+
+  EXPECT_TRUE(map.erase(key2));
+  EXPECT_FALSE(map.erase(key2));
+
+  {
+    auto range = map.getRange();
+    EXPECT_EQ(std::distance(range.begin(), range.end()), 0);
+  }
+}