[orc-rt] Add MemProt, MemLifetime, AllocGroup, and AllocGroupMap. (#157078)

MemProt and MemLifetime are enum classes representing memory protection
(Read | Write | Exec) and lifetime policy (Standard or Finalize-only)
respectively.

An AllocGroup is a compressed (MemProt, MemLifetime) pair.

AllocGroupSmallMap<T> is a compressed map of AllocGroup -> T.

These utilities will be used in upcoming memory management APIs in the
ORC runtime.

Author: lhames

orc-rt/include/CMakeLists.txt

@@ -9,6 +9,7 @@ set(ORC_RT_HEADERS
     orc-rt/IntervalMap.h
     orc-rt/IntervalSet.h
     orc-rt/Math.h
+    orc-rt/MemoryFlags.h
     orc-rt/RTTI.h
     orc-rt/WrapperFunction.h
     orc-rt/SimplePackedSerialization.h

orc-rt/include/orc-rt/BitmaskEnum.h

@@ -17,6 +17,7 @@
 #define ORC_RT_BITMASKENUM_H
 
 #include "Math.h"
+#include "bit.h"
 
 #include <cassert>
 #include <type_traits>
@@ -114,6 +115,14 @@ constexpr std::underlying_type_t<E> bitmask_enum_to_underlying(E Val) noexcept {
   return U;
 }
 
+template <typename E, typename _ = std::enable_if_t<is_bitmask_enum_v<E>>>
+struct bitmask_enum_num_bits {
+  static constexpr int value = bit_width(largest_bitmask_enum_bit<E>::value);
+};
+
+template <typename E>
+inline constexpr int bitmask_enum_num_bits_v = bitmask_enum_num_bits<E>::value;
+
 template <typename E, typename = std::enable_if_t<is_bitmask_enum_v<E>>>
 constexpr E operator~(E Val) noexcept {
   return static_cast<E>(~bitmask_enum_to_underlying(Val) &

orc-rt/include/orc-rt/MemoryFlags.h

@@ -0,0 +1,136 @@
+//===--------- MemoryFlags.h -- Memory allocation flags ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Memory allocation flags.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_MEMORYFLAGS_H
+#define ORC_RT_MEMORYFLAGS_H
+
+#include "orc-rt/BitmaskEnum.h"
+#include "orc-rt/bit.h"
+
+#include <algorithm>
+#include <utility>
+#include <vector>
+
+namespace orc_rt {
+
+/// Describes Read/Write/Exec permissions for memory.
+enum class MemProt : unsigned {
+  None = 0,
+  Read = 1U << 0,
+  Write = 1U << 1,
+  Exec = 1U << 2,
+  ORC_RT_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Exec)
+};
+
+/// Describes a memory lifetime policy.
+enum class MemLifetime : unsigned {
+  /// Standard memory should be deallocated by the corresponding call to
+  /// deallocate.
+  Standard,
+
+  /// Finalize memory should be deallocated at the end of the finalization
+  /// process.
+  Finalize
+};
+
+/// A pair of memory protections and lifetime policy.
+class AllocGroup {
+private:
+  static constexpr int NumProtBits = bitmask_enum_num_bits_v<MemProt>;
+  static constexpr int NumLifetimeBits = 1;
+  static constexpr int NumBits = NumProtBits + NumLifetimeBits;
+
+  typedef uint8_t underlying_type;
+
+  static_assert(NumBits <= std::numeric_limits<underlying_type>::digits,
+                "Not enough bits to hold (prot, lifetime) pair");
+
+  constexpr static underlying_type ProtMask = (1U << NumProtBits) - 1;
+  constexpr static underlying_type LifetimeMask = (1U << NumLifetimeBits) - 1;
+
+public:
+  static constexpr size_t MaxValues = 1U << NumBits;
+
+  AllocGroup() = default;
+  AllocGroup(MemProt MP, MemLifetime ML = MemLifetime::Standard)
+      : Id((static_cast<underlying_type>(ML) << NumProtBits) |
+           static_cast<underlying_type>(MP)) {}
+
+  MemProt getMemProt() const { return static_cast<MemProt>(Id & ProtMask); }
+
+  MemLifetime getMemLifetime() const {
+    return static_cast<MemLifetime>((Id >> NumProtBits) & LifetimeMask);
+  }
+
+  friend bool operator==(const AllocGroup &LHS, const AllocGroup &RHS) {
+    return LHS.Id == RHS.Id;
+  }
+
+  friend bool operator!=(const AllocGroup &LHS, const AllocGroup &RHS) {
+    return !(LHS == RHS);
+  }
+
+  friend bool operator<(const AllocGroup &LHS, const AllocGroup &RHS) {
+    return LHS.Id < RHS.Id;
+  }
+
+private:
+  underlying_type Id = 0;
+};
+
+/// A specialized small-map for AllocGroups.
+///
+/// Iteration order is guaranteed to match key ordering.
+template <typename T> class AllocGroupSmallMap {
+private:
+  using ElemT = std::pair<AllocGroup, T>;
+  using VectorTy = std::vector<ElemT>;
+
+  static bool compareKey(const ElemT &E, const AllocGroup &G) {
+    return E.first < G;
+  }
+
+public:
+  using iterator = typename VectorTy::iterator;
+
+  AllocGroupSmallMap() = default;
+  AllocGroupSmallMap(std::initializer_list<std::pair<AllocGroup, T>> Inits)
+      : Elems(Inits) {
+    std::sort(Elems, [](const ElemT &LHS, const ElemT &RHS) {
+      return LHS.first < RHS.first;
+    });
+  }
+
+  iterator begin() { return Elems.begin(); }
+  iterator end() { return Elems.end(); }
+  iterator find(AllocGroup G) {
+    auto I = std::lower_bound(Elems.begin(), Elems.end(), G, compareKey);
+    return (I == end() || I->first == G) ? I : end();
+  }
+
+  bool empty() const { return Elems.empty(); }
+  size_t size() const { return Elems.size(); }
+
+  T &operator[](AllocGroup G) {
+    auto I = std::lower_bound(Elems.begin(), Elems.end(), G, compareKey);
+    if (I == Elems.end() || I->first != G)
+      I = Elems.insert(I, std::make_pair(G, T()));
+    return I->second;
+  }
+
+private:
+  VectorTy Elems;
+};
+
+} // namespace orc_rt
+
+#endif // ORC_RT_MEMORYFLAGS_H

orc-rt/unittests/CMakeLists.txt

@@ -19,6 +19,7 @@ add_orc_rt_unittest(CoreTests
   IntervalMapTest.cpp
   IntervalSetTest.cpp
   MathTest.cpp
+  MemoryFlagsTest.cpp
   RTTITest.cpp
   SimplePackedSerializationTest.cpp
   WrapperFunctionBufferTest.cpp

orc-rt/unittests/MemoryFlagsTest.cpp

@@ -0,0 +1,89 @@
+//===- MemoryFlags.cpp ----------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for orc-rt's MemoryFlags.h APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/MemoryFlags.h"
+#include "gtest/gtest.h"
+
+using namespace orc_rt;
+
+TEST(MemProtTest, Basics) {
+  MemProt MPNone = MemProt::None;
+
+  EXPECT_EQ(MPNone & MemProt::Read, MemProt::None);
+  EXPECT_EQ(MPNone & MemProt::Write, MemProt::None);
+  EXPECT_EQ(MPNone & MemProt::Exec, MemProt::None);
+
+  EXPECT_EQ(MPNone | MemProt::Read, MemProt::Read);
+  EXPECT_EQ(MPNone | MemProt::Write, MemProt::Write);
+  EXPECT_EQ(MPNone | MemProt::Exec, MemProt::Exec);
+
+  MemProt MPAll = MemProt::Read | MemProt::Write | MemProt::Exec;
+  EXPECT_EQ(MPAll & MemProt::Read, MemProt::Read);
+  EXPECT_EQ(MPAll & MemProt::Write, MemProt::Write);
+  EXPECT_EQ(MPAll & MemProt::Exec, MemProt::Exec);
+}
+
+TEST(AllocGroupTest, Default) {
+  AllocGroup G;
+  EXPECT_EQ(G.getMemProt(), MemProt::None);
+  EXPECT_EQ(G.getMemLifetime(), MemLifetime::Standard);
+}
+
+TEST(AllocGroupTest, InitMemProtOnly) {
+  AllocGroup G(MemProt::Read | MemProt::Write);
+  EXPECT_EQ(G.getMemProt(), MemProt::Read | MemProt::Write);
+  EXPECT_EQ(G.getMemLifetime(), MemLifetime::Standard);
+}
+
+TEST(AllocGroupTest, InitMemProtAndLifetime) {
+  AllocGroup G(MemProt::Read | MemProt::Write, MemLifetime::Finalize);
+  EXPECT_EQ(G.getMemProt(), MemProt::Read | MemProt::Write);
+  EXPECT_EQ(G.getMemLifetime(), MemLifetime::Finalize);
+}
+
+TEST(AllocGroupTest, Equality) {
+  AllocGroup G(MemProt::Read, MemLifetime::Standard);
+  EXPECT_EQ(G, AllocGroup(MemProt::Read, MemLifetime::Standard));
+  EXPECT_NE(G, AllocGroup(MemProt::Write, MemLifetime::Standard));
+  EXPECT_NE(G, AllocGroup(MemProt::Read, MemLifetime::Finalize));
+}
+
+TEST(AllocGroupTest, LessThan) {
+  // Check that AllocGroups can be compared via less-than so that they can be
+  // used as keys in standard containers.
+  EXPECT_LT(AllocGroup(MemProt::Read, MemLifetime::Standard),
+            AllocGroup(MemProt::Write, MemLifetime::Standard));
+  EXPECT_LT(AllocGroup(MemProt::Exec, MemLifetime::Standard),
+            AllocGroup(MemProt::Read, MemLifetime::Finalize));
+}
+
+TEST(AllocGroupSmallMap, EmptyMap) {
+  AllocGroupSmallMap<bool> EM;
+  EXPECT_TRUE(EM.empty());
+  EXPECT_EQ(EM.size(), 0u);
+}
+
+TEST(AllocGroupSmallMap, NonEmptyMap) {
+  AllocGroupSmallMap<unsigned> NEM;
+  NEM[MemProt::Read] = 42;
+
+  EXPECT_FALSE(NEM.empty());
+  EXPECT_EQ(NEM.size(), 1U);
+  EXPECT_EQ(NEM[MemProt::Read], 42U);
+  EXPECT_EQ(NEM.find(MemProt::Read), NEM.begin());
+  EXPECT_EQ(NEM.find(MemProt::Read | MemProt::Write), NEM.end());
+
+  NEM[MemProt::Read | MemProt::Write] = 7;
+  EXPECT_EQ(NEM.size(), 2U);
+  EXPECT_EQ(NEM.begin()->second, 42U);
+  EXPECT_EQ((NEM.begin() + 1)->second, 7U);
+}