[basic] Add a simple vector backed 2 stage multi map.

I have been using this in a bunch of places in the compiler and rather than
implement it by hand over and over (and maybe messing up), this commit just
commits a correct implementation.

This data structure is a map backed by a vector like data structure. It has two
phases:

1. An insertion phase when the map is mutable and one inserts (key, value) pairs
into the map. These are just appeneded into the storage array.

2. A frozen stage when the map is immutable and one can now perform map queries
on the multimap.

The map transitions from the mutable, thawed phase to the immutable, frozen
phase by performing a stable_sort of its internal storage by only the key. Since
this is a stable_sort, we know that the relative insertion order of values is
preserved if their keys equal. Thus the sorting will have created contiguous
regions in the array of values, all mapped to the same key, that are insertion
order. Thus by finding the lower_bound for a given key, we are guaranteed to get
the first element in that continguous range. We can then do a forward search to
find the end of the region, allowing us to then return an ArrayRef to these
internal values.

The reason why I keep on finding myself using this is that this map enables one
to map a key to an array of values without needing to store small vectors in a
map or use heap allocated memory, all key, value pairs are stored inline (in
potentially a single SmallVector given that one is using SmallFrozenMultiMap).

include/swift/Basic/FrozenMultiMap.h

@@ -0,0 +1,187 @@
+//===--- FrozenMultiMap.h ----------------------------------*- C++ --------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+///
+/// \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.
+///
+/// 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.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_FROZENMULTIMAP_H
+#define SWIFT_BASIC_FROZENMULTIMAP_H
+
+#include "swift/Basic/LLVM.h"
+#include "swift/Basic/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include <vector>
+
+namespace swift {
+
+template <typename Key, typename Value,
+          typename VectorStorage = std::vector<std::pair<Key, Value>>>
+class FrozenMultiMap {
+  VectorStorage storage;
+  bool frozen = false;
+
+private:
+  struct PairToSecondElt;
+
+public:
+  using PairToSecondEltRange =
+      TransformRange<ArrayRef<std::pair<Key, Value>>, PairToSecondElt>;
+
+  FrozenMultiMap() = default;
+
+  void insert(const Key &key, const Value &value) {
+    assert(!isFrozen() && "Can not insert new keys once map is frozen");
+    storage.emplace_back(key, value);
+  }
+
+  Optional<PairToSecondEltRange> find(const Key &key) const {
+    assert(isFrozen() &&
+           "Can not perform a find 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, Value> &p1, const std::pair<Key, Value> &p2) {
+          return p1.first < p2.first;
+        });
+    if (start == storage.end() || start->first != key) {
+      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; });
+    unsigned count = std::distance(start, end);
+    ArrayRef<std::pair<Key, Value>> slice(&*start, count);
+    return PairToSecondEltRange(slice, PairToSecondElt());
+  }
+
+  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) {
+                       // Only compare the first entry so that we preserve
+                       // insertion order.
+                       return lhs.first < rhs.first;
+                     });
+    frozen = true;
+  }
+
+  unsigned size() const { return storage.size(); }
+  bool empty() const { return storage.empty(); }
+
+  struct iterator : std::iterator<std::forward_iterator_tag,
+                                  std::pair<Key, ArrayRef<Value>>> {
+    using base_iterator = typename decltype(storage)::iterator;
+
+    FrozenMultiMap &map;
+    base_iterator baseIter;
+    Optional<std::pair<Key, PairToSecondEltRange>> currentValue;
+
+    iterator(FrozenMultiMap &map, base_iterator iter)
+        : map(map), baseIter(iter), currentValue() {
+      // If we are end, just return.
+      if (iter == map.storage.end()) {
+        return;
+      }
+
+      // Otherwise, prime our first range.
+      updateCurrentValue();
+    }
+
+    void updateCurrentValue() {
+      base_iterator end = map.storage.end();
+      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())};
+    }
+
+    iterator &operator++() {
+      baseIter = std::find_if_not(std::next(baseIter), map.storage.end(),
+                                  [&](const std::pair<Key, Value> &elt) {
+                                    return elt.first == baseIter->first;
+                                  });
+      updateCurrentValue();
+      return *this;
+    }
+
+    iterator operator++(int) {
+      auto tmp = *this;
+      baseIter = std::find_if_not(std::next(baseIter), map.storage.end(),
+                                  [&](const std::pair<Key, Value> &elt) {
+                                    return elt.first == baseIter->first;
+                                  });
+      updateCurrentValue();
+      return tmp;
+    }
+
+    std::pair<Key, PairToSecondEltRange> operator*() const {
+      return *currentValue;
+    }
+
+    bool operator==(const iterator &RHS) const {
+      return baseIter == RHS.baseIter;
+    }
+
+    bool operator!=(const iterator &RHS) const {
+      return baseIter != RHS.baseIter;
+    }
+  };
+
+  /// 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.
+  llvm::iterator_range<iterator> getRange() const {
+    assert(isFrozen() &&
+           "Can not create range until data structure is frozen?!");
+    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);
+  }
+};
+
+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;
+  }
+};
+
+template <typename Key, typename Value, unsigned SmallSize>
+using SmallFrozenMultiMap =
+    FrozenMultiMap<Key, Value, SmallVector<std::pair<Key, Value>, SmallSize>>;
+
+} // namespace swift
+
+#endif

unittests/Basic/CMakeLists.txt

@@ -15,6 +15,7 @@ add_swift_unittest(SwiftBasicTests
   EncodedSequenceTest.cpp
   ExponentialGrowthAppendingBinaryByteStreamTests.cpp
   FileSystemTest.cpp
+  FrozenMultiMapTest.cpp
   ImmutablePointerSetTest.cpp
   JSONSerialization.cpp
   OptionSetTest.cpp

unittests/Basic/FrozenMultiMapTest.cpp

@@ -0,0 +1,187 @@
+//===--- FrozenMultiMapTest.cpp -------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#define DEBUG_TYPE "swift-frozen-multi-map-test"
+
+#include "swift/Basic/FrozenMultiMap.h"
+#include "swift/Basic/LLVM.h"
+#include "swift/Basic/Lazy.h"
+#include "swift/Basic/NullablePtr.h"
+#include "swift/Basic/Range.h"
+#include "swift/Basic/STLExtras.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "gtest/gtest.h"
+#include <map>
+#include <random>
+#include <set>
+
+using namespace swift;
+
+namespace {
+
+class Canary {
+  static unsigned currentID;
+  unsigned id;
+
+public:
+  static void resetIDs() { currentID = 0; }
+  Canary(unsigned id) : id(id) {}
+  Canary() {
+    id = currentID;
+    ++currentID;
+  }
+
+  unsigned getID() const { return id; }
+  bool operator<(const Canary &other) const { return id < other.id; }
+
+  bool operator==(const Canary &other) const { return id == other.id; }
+
+  bool operator!=(const Canary &other) const { return !(*this == other); }
+};
+
+unsigned Canary::currentID = 0;
+
+} // namespace
+
+TEST(FrozenMultiMapCustomTest, SimpleFind) {
+  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);
+  {
+    auto r = map.find(key1);
+    EXPECT_TRUE(r.hasValue());
+    EXPECT_EQ(r->size(), 3u);
+    EXPECT_EQ((*r)[0].getID(), 2u);
+    EXPECT_EQ((*r)[1].getID(), 3u);
+    EXPECT_EQ((*r)[2].getID(), 4u);
+  }
+
+  {
+    auto r = map.find(key2);
+    EXPECT_TRUE(r.hasValue());
+    EXPECT_EQ(r->size(), 2u);
+    EXPECT_EQ((*r)[0].getID(), 5u);
+    EXPECT_EQ((*r)[1].getID(), 6u);
+  }
+}
+
+TEST(FrozenMultiMapCustomTest, SimpleIter) {
+  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);
+
+  auto range = map.getRange();
+
+  EXPECT_EQ(std::distance(range.begin(), range.end()), 2);
+
+  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);
+  }
+
+  ++iter;
+  {
+    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);
+  }
+}
+
+TEST(FrozenMultiMapCustomTest, RandomAgainstStdMultiMap) {
+  Canary::resetIDs();
+  FrozenMultiMap<unsigned, unsigned> map;
+  std::multimap<unsigned, unsigned> stdMultiMap;
+
+  auto seed =
+      std::chrono::high_resolution_clock::now().time_since_epoch().count();
+  std::mt19937 mt_rand(seed);
+
+  std::vector<unsigned> keyIdList;
+  for (unsigned i = 0; i < 1024; ++i) {
+    unsigned keyID = mt_rand() % 20;
+    keyIdList.push_back(keyID);
+    for (unsigned valueID = (mt_rand()) % 15; valueID < 15; ++valueID) {
+      map.insert(keyID, valueID);
+      stdMultiMap.emplace(keyID, valueID);
+    }
+  }
+
+  map.setFrozen();
+
+  // Then for each key.
+  for (unsigned i : keyIdList) {
+    // Make sure that we have the same elements in the same order for each key.
+    auto range = *map.find(i);
+    auto stdRange = stdMultiMap.equal_range(i);
+    EXPECT_EQ(std::distance(range.begin(), range.end()),
+              std::distance(stdRange.first, stdRange.second));
+    auto modernStdRange = llvm::make_range(stdRange.first, stdRange.second);
+    for (auto p : llvm::zip(range, modernStdRange)) {
+      unsigned lhs = std::get<0>(p);
+      unsigned rhs = std::get<1>(p).second;
+      EXPECT_EQ(lhs, rhs);
+    }
+  }
+
+  // Then check that when we iterate over both ranges, we get the same order.
+  {
+    auto range = map.getRange();
+    auto rangeIter = range.begin();
+    auto stdRangeIter = stdMultiMap.begin();
+
+    while (rangeIter != range.end()) {
+      auto rangeElt = *rangeIter;
+
+      for (unsigned i : indices(rangeElt.second)) {
+        EXPECT_EQ(rangeElt.first, stdRangeIter->first);
+        EXPECT_EQ(rangeElt.second[i], stdRangeIter->second);
+        ++stdRangeIter;
+      }
+
+      ++rangeIter;
+    }
+  }
+}