Add some optimized data structures for maps on a small number of

keys, especially enumerated keys.

Author: rjmccall

include/swift/Basic/EnumMap.h

@@ -0,0 +1,203 @@
+//===--- EnumMap.h - A map optimized for having enum keys -------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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
+//
+//===----------------------------------------------------------------------===//
+///
+///  This file defines the EnumMap class template, which is a map data
+///  structure optimized for working with enumerated keys.  It is built on
+///  top of SmallMap, but it replaces the default large map with a flat
+///  heap-allocated array of indexes into the elements, which is reasonable
+///  for small-ish enums.
+///
+///  Currently the map requires the key type to be an enum type.
+///  The expectation is that the enum has a small number of enumerators
+///  which are all in the range 0..<NumValues.  NumValues must be provided
+///  by specializing the EnumTraits class.
+///
+///  The elements of the map remain insertion-ordered for the lifetime of
+///  the map.  There are currently no operations to remove elements.
+///  Iterators are invalidated by insertion.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_ENUMMAP_H
+#define SWIFT_BASIC_ENUMMAP_H
+
+#include "swift/Basic/EnumTraits.h"
+#include "swift/Basic/SmallMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include <type_traits>
+
+namespace swift {
+
+/// The maximum number of elements that the map can have before
+/// it flips from brute-force searching the keys to using a sparse
+/// array.
+static constexpr size_t DefaultEnumMapDirectSearchLimit =
+  DefaultSmallMapDirectSearchLimit;
+
+/// The primary customization point for an EnumMap.
+///
+/// template <>
+/// struct EnumMapTraits<MyKey> {
+///   using IndexType = <some integer type>;
+///   struct LargeMapStorage {
+///     std::optional<IndexType> find(IndexType) const;
+///     std::pair<IndexType, bool> insert(IndexType key, IndexType value);
+///   };
+/// };
+template <class Key, class KeyTraits = EnumTraits<Key>>
+struct EnumMapTraits;
+
+template <class Key, class Value,
+          size_t DirectSearchLimit = DefaultEnumMapDirectSearchLimit,
+          class MapTraits = EnumMapTraits<Key>,
+          class ElementStorage = llvm::SmallVector<Value>>
+class EnumMap {
+  using IndexType = typename MapTraits::IndexType;
+
+  // EnumMapTraits is currently designed to be usable directly as a
+  // SmallMapTraits.
+  using MapType =
+    SmallMap<IndexType, Value, DirectSearchLimit, MapTraits, ElementStorage>;
+  MapType map;
+
+public:
+  bool empty() const { return map.empty(); }
+  size_t size() const { return map.size(); }
+
+  using iterator = typename MapType::iterator;
+  iterator begin() { return map.begin(); }
+  iterator end() { return map.end(); }
+
+  using const_iterator = typename MapType::const_iterator;
+  const_iterator begin() const { return map.begin(); }
+  const_iterator end() const { return map.end(); }
+
+  /// Look up a key in the map.  Returns end() if the entry is not found.
+  const_iterator find(Key key) const {
+    return map.find(IndexType(key));
+  }
+
+  /// Try to insert the given key/value pair.  If there's already an element
+  /// with this key, return false and an iterator for the existing element.
+  /// Otherwise, return true and an iterator for the new element.
+  ///
+  /// The value in the set will be constructed by emplacing it with the
+  /// given arguments.
+  template <class... Args>
+  std::pair<iterator, bool> insert(Key key, Args &&...valueArgs) {
+    return map.insert(IndexType(key), std::forward<Args>(valueArgs)...);
+  }
+};
+
+namespace EnumMapImpl {
+
+template <size_t N,
+          bool SmallEnoughForUInt8 = (N < (1U << 8)),
+          bool SmallEnoughForUInt16 = (N < (1U << 16))>
+struct SufficientIntFor;
+
+template <size_t N>
+struct SufficientIntFor<N, true, true> {
+  using type = uint8_t;
+};
+
+template <size_t N>
+struct SufficientIntFor<N, false, true> {
+  using type = uint16_t;
+};
+
+template <size_t N>
+struct SufficientIntFor<N, false, false> {
+  static_assert(N < (1ULL << 32), "just how large is this \"enum\" exactly");
+  using type = uint32_t;
+};
+
+/// A map from integers in 0..<N to integers in 0..<N, implemented as a
+/// flat array of integers in 0...N, with zero meaning a missing entry.
+///
+/// This is a great implementation for N <= 255, where the
+/// entire flat array is <= 256 bytes.  It gets increasingly marginal
+/// for N up to ~1K or so (unless we really expect to have entries
+/// for a large proportion of the enum).  Past that, we should probably
+/// be falling back on something like a hashtable, because needing tens
+/// of kilobytes to hold as few as 17 entries is objectively unreasonable.
+template <size_t N>
+class FlatMap {
+public:
+  using IndexType = typename SufficientIntFor<N>::type;
+  using StoredIndexType = typename SufficientIntFor<N + 1>::type;
+
+private:
+  StoredIndexType *ptr;
+
+public:
+  FlatMap() : ptr(new StoredIndexType[N]) {
+    memset(ptr, 0, N * sizeof(StoredIndexType));
+  }
+  FlatMap(FlatMap &&other)
+      : ptr(other.ptr) {
+    other.ptr = nullptr;
+  }
+  FlatMap &operator=(FlatMap &&other) {
+    delete ptr;
+    ptr = other.ptr;
+    other.ptr = nullptr;
+  }
+  FlatMap(const FlatMap &other)
+      : ptr(new StoredIndexType[N]) {
+    memcpy(ptr, other.ptr, N * sizeof(StoredIndexType));
+  }
+  FlatMap &operator=(const FlatMap &other) {
+    memcpy(ptr, other.ptr, N * sizeof(StoredIndexType));
+  }
+
+  ~FlatMap() {
+    delete ptr;
+  }
+
+  std::pair<IndexType, bool> insert(IndexType key, IndexType value) {
+    assert(key < N);
+    StoredIndexType &entry = ptr[key];
+    if (entry == 0) {
+      entry = StoredIndexType(value) + 1;
+      return std::make_pair(value, true);
+    } else {
+      return std::make_pair(IndexType(entry - 1), false);
+    }
+  }
+
+  std::optional<IndexType> find(IndexType key) const {
+    assert(key < N);
+    StoredIndexType entry = ptr[key];
+    if (entry == 0) {
+      return std::nullopt;
+    } else {
+      return IndexType(entry - 1);
+    }
+  };
+};
+
+} // end namespace EnumMapImpl
+
+/// The default implementation of EnumMapTraits.
+template <class Key_, class KeyTraits_>
+struct EnumMapTraits {
+  using Key = Key_;
+  using KeyTraits = KeyTraits_;
+
+  using LargeMapStorage = EnumMapImpl::FlatMap<KeyTraits::NumValues>;
+  using IndexType = typename LargeMapStorage::IndexType;
+};
+
+} // end namespace swift
+
+#endif

include/swift/Basic/EnumTraits.h

@@ -0,0 +1,34 @@
+//===--- EnumTraits.h - Traits for densely-packed enums ---------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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
+//
+//===----------------------------------------------------------------------===//
+///
+///  This file defines the EnumTraits concept, which can be used to
+///  communicate information about an enum type's enumerators that currently
+///  can't be recovered from the compiler.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_ENUMTRAITS_H
+#define SWIFT_BASIC_ENUMTRAITS_H
+
+namespace swift {
+
+/// A simple traits concept for recording the number of cases in an enum.
+///
+///  template <> class EnumTraits<WdigetKind> {
+///    static constexpr size_t NumValues = NumWidgetKinds;
+///  };
+template <class E>
+struct EnumTraits;
+
+} // end namespace swift
+
+#endif