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

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;
+    }
+  }
+}