| 1 | /*++ |
| 2 | |
| 3 | Copyright (c) Microsoft. All rights reserved. |
| 4 | |
| 5 | Module Name: |
| 6 | |
| 7 | EnumVariantMap.h |
| 8 | |
| 9 | Abstract: |
| 10 | |
| 11 | Template for enum-based variant maps. |
| 12 | |
| 13 | --*/ |
| 14 | #pragma once |
| 15 | #include <map> |
| 16 | #include <type_traits> |
| 17 | #include <utility> |
| 18 | #include <variant> |
| 19 | #include <vector> |
| 20 | |
| 21 | // This template set is used for Arg storage and Context Data storage by enum type. |
| 22 | // The backing storage is a std::multimap of the enum to a variant of types. |
| 23 | // This enables strongly typed storage and retrieval of values based on an enum key. |
| 24 | namespace wsl::windows::wslc { |
| 25 | |
| 26 | // Enum based variant helper. |
| 27 | // Enum must be an enum whose first member has the value 0, each subsequent member increases by 1, and the final member is named Max. |
| 28 | // Mapping is a template type that takes one template parameter of type Enum, and whose members define value_t as the type for that enum value. |
| 29 | template <typename Enum, template <Enum> typename Mapping> |
| 30 | struct EnumBasedVariant |
| 31 | { |
| 32 | private: |
| 33 | // Used to deduce the variant type; making a variant that includes std::monostate and all Mapping types. |
| 34 | template <size_t... I> |
| 35 | static inline auto Deduce(std::index_sequence<I...>) |
| 36 | { |
| 37 | return std::variant<std::monostate, typename Mapping<static_cast<Enum>(I)>::value_t...>{}; |
| 38 | } |
| 39 | |
| 40 | public: |
| 41 | // Holds data of any type listed in Mapping. |
| 42 | using variant_t = decltype(Deduce(std::make_index_sequence<static_cast<size_t>(Enum::Max)>())); |
| 43 | |
| 44 | // Gets the index into the variant for the given Data. |
| 45 | static constexpr inline size_t Index(Enum e) |
| 46 | { |
| 47 | return static_cast<size_t>(e) + 1; |
| 48 | } |
| 49 | }; |
| 50 | |
| 51 | // An action that can be taken on an EnumBasedVariantMap. |
| 52 | enum class EnumBasedVariantMapAction |
| 53 | { |
| 54 | Add, |
| 55 | Contains, |
| 56 | Get, |
| 57 | GetMutable, |
| 58 | GetAll, |
| 59 | Count, |
| 60 | Remove, |
| 61 | }; |
| 62 | |
| 63 | // A callback function that can take any action in response to map operations, such as logging |
| 64 | // accesses or maintaining state derived from the map contents. |
| 65 | template <typename Enum> |
| 66 | using EnumBasedVariantMapActionCallback = void (*)(const void* map, Enum value, EnumBasedVariantMapAction action); |
| 67 | |
| 68 | // Forward declaration for EnumBasedVariantMapEmplacer |
| 69 | template <typename Enum, template <Enum> typename Mapping, typename V> |
| 70 | struct EnumBasedVariantMapEmplacer; |
| 71 | |
| 72 | // Provides a multimap of the Enum to the mapped types (allows multiple values per key). |
| 73 | template <typename Enum, template <Enum> typename Mapping, EnumBasedVariantMapActionCallback<Enum> Callback = nullptr> |
| 74 | struct EnumBasedVariantMap |
| 75 | { |
| 76 | using Variant = EnumBasedVariant<Enum, Mapping>; |
| 77 | |
| 78 | template <Enum E> |
| 79 | using mapping_t = typename Mapping<E>::value_t; |
| 80 | |
| 81 | // Adds a value to the map. With multimap, this always adds a new entry (doesn't overwrite). |
| 82 | template <Enum E> |
| 83 | void Add(mapping_t<E>&& v) |
| 84 | { |
| 85 | if constexpr (Callback) |
| 86 | { |
| 87 | Callback(this, E, EnumBasedVariantMapAction::Add); |
| 88 | } |
| 89 | |
| 90 | // Compile-time type checking - this should always pass since mapping_t<E> is the correct type |
| 91 | using CleanV = std::remove_cvref_t<mapping_t<E>>; |
| 92 | static_assert( |
| 93 | std::is_same_v<CleanV, mapping_t<E>>, |
| 94 | "Type mismatch in Add: provided type does not match the expected type for this enum value"); |
| 95 | |
| 96 | typename Variant::variant_t variant; |
| 97 | variant.template emplace<Variant::Index(E)>(std::move(v)); |
| 98 | m_data.emplace(E, std::move(variant)); |
| 99 | } |
| 100 | |
| 101 | template <Enum E> |
| 102 | void Add(const mapping_t<E>& v) |
| 103 | { |
| 104 | if constexpr (Callback) |
| 105 | { |
| 106 | Callback(this, E, EnumBasedVariantMapAction::Add); |
| 107 | } |
| 108 | |
| 109 | // Compile-time type checking - this should always pass since mapping_t<E> is the correct type |
| 110 | using CleanV = std::remove_cvref_t<mapping_t<E>>; |
| 111 | static_assert( |
| 112 | std::is_same_v<CleanV, mapping_t<E>>, |
| 113 | "Type mismatch in Add: provided type does not match the expected type for this enum value"); |
| 114 | |
| 115 | typename Variant::variant_t variant; |
| 116 | variant.template emplace<Variant::Index(E)>(v); |
| 117 | m_data.emplace(E, std::move(variant)); |
| 118 | } |
| 119 | |
| 120 | // Runtime version of Add that takes the enum as a parameter. |
| 121 | template <typename V> |
| 122 | void Add(Enum e, V&& v) |
| 123 | { |
| 124 | if constexpr (Callback) |
| 125 | { |
| 126 | Callback(this, e, EnumBasedVariantMapAction::Add); |
| 127 | } |
| 128 | |
| 129 | // Check if the type matches the SPECIFIC enum value at compile time if possible |
| 130 | using CleanV = std::remove_cvref_t<V>; |
| 131 | |
| 132 | // Pre-check if this type matches the specific enum value being added to |
| 133 | if (!IsMatchingType<CleanV>(e)) |
| 134 | { |
| 135 | THROW_HR_MSG(E_INVALIDARG, "Type mismatch: provided type does not match the expected type for enum value %d", static_cast<int>(e)); |
| 136 | } |
| 137 | |
| 138 | typename Variant::variant_t variant; |
| 139 | EmplaceAtRuntimeIndex(variant, e, std::forward<V>(v), std::make_index_sequence<static_cast<size_t>(Enum::Max)>()); |
| 140 | m_data.emplace(e, std::move(variant)); |
| 141 | } |
| 142 | |
| 143 | // Runtime method to check if value V matches the mapped type for an enum value. |
| 144 | template <typename V> |
| 145 | bool IsMatchingType(Enum e) const |
| 146 | { |
| 147 | return IsMatchingTypeImpl<V>(e, std::make_index_sequence<static_cast<size_t>(Enum::Max)>()); |
| 148 | } |
| 149 | |
| 150 | // Return a value indicating whether the given enum has at least one entry. |
| 151 | bool Contains(Enum e) const |
| 152 | { |
| 153 | if constexpr (Callback) |
| 154 | { |
| 155 | Callback(this, e, EnumBasedVariantMapAction::Contains); |
| 156 | } |
| 157 | return (m_data.find(e) != m_data.end()); |
| 158 | } |
| 159 | |
| 160 | // Gets the count of values for a specific enum key. |
| 161 | size_t Count(Enum e) const |
| 162 | { |
| 163 | if constexpr (Callback) |
| 164 | { |
| 165 | Callback(this, e, EnumBasedVariantMapAction::Count); |
| 166 | } |
| 167 | return m_data.count(e); |
| 168 | } |
| 169 | |
| 170 | // Gets the FIRST value for the enum key (for backward compatibility). |
| 171 | // Non-const version returns a reference that can be modified. |
| 172 | template <Enum E> |
| 173 | mapping_t<E>& Get() |
| 174 | { |
| 175 | if constexpr (Callback) |
| 176 | { |
| 177 | Callback(this, E, EnumBasedVariantMapAction::GetMutable); |
| 178 | } |
| 179 | auto itr = m_data.find(E); |
| 180 | THROW_HR_IF_MSG(E_NOT_SET, itr == m_data.end(), "Get(%d): key not found", static_cast<int>(E)); |
| 181 | |
| 182 | // Validate that the variant holds the expected type at the expected index |
| 183 | constexpr size_t expectedIndex = Variant::Index(E); |
| 184 | if (itr->second.index() != expectedIndex) |
| 185 | { |
| 186 | THROW_HR_MSG( |
| 187 | E_UNEXPECTED, |
| 188 | "Get(%d): variant type mismatch - expected index %zu, got %zu", |
| 189 | static_cast<int>(E), |
| 190 | expectedIndex, |
| 191 | itr->second.index()); |
| 192 | } |
| 193 | |
| 194 | return std::get<expectedIndex>(itr->second); |
| 195 | } |
| 196 | |
| 197 | // Const overload of Get, cannot be modified. |
| 198 | template <Enum E> |
| 199 | const mapping_t<E>& Get() const |
| 200 | { |
| 201 | if constexpr (Callback) |
| 202 | { |
| 203 | Callback(this, E, EnumBasedVariantMapAction::Get); |
| 204 | } |
| 205 | auto itr = m_data.find(E); |
| 206 | THROW_HR_IF_MSG(E_NOT_SET, itr == m_data.cend(), "Get(%d): key not found", static_cast<int>(E)); |
| 207 | |
| 208 | // Validate that the variant holds the expected type at the expected index |
| 209 | constexpr size_t expectedIndex = Variant::Index(E); |
| 210 | if (itr->second.index() != expectedIndex) |
| 211 | { |
| 212 | THROW_HR_MSG( |
| 213 | E_UNEXPECTED, |
| 214 | "Get(%d): variant type mismatch - expected index %zu, got %zu", |
| 215 | static_cast<int>(E), |
| 216 | expectedIndex, |
| 217 | itr->second.index()); |
| 218 | } |
| 219 | |
| 220 | return std::get<expectedIndex>(itr->second); |
| 221 | } |
| 222 | |
| 223 | // Gets ALL values for a specific enum key as a vector. |
| 224 | template <Enum E> |
| 225 | std::vector<mapping_t<E>> GetAll() const |
| 226 | { |
| 227 | if constexpr (Callback) |
| 228 | { |
| 229 | Callback(this, E, EnumBasedVariantMapAction::GetAll); |
| 230 | } |
| 231 | |
| 232 | std::vector<mapping_t<E>> results; |
| 233 | auto range = m_data.equal_range(E); |
| 234 | |
| 235 | for (auto it = range.first; it != range.second; ++it) |
| 236 | { |
| 237 | results.push_back(std::get<Variant::Index(E)>(it->second)); |
| 238 | } |
| 239 | |
| 240 | return results; |
| 241 | } |
| 242 | |
| 243 | // Removes ALL entries for a specific enum key. |
| 244 | void Remove(Enum e) |
| 245 | { |
| 246 | if constexpr (Callback) |
| 247 | { |
| 248 | Callback(this, e, EnumBasedVariantMapAction::Remove); |
| 249 | } |
| 250 | m_data.erase(e); |
| 251 | } |
| 252 | |
| 253 | // Gets the total number of items stored (across all keys). |
| 254 | size_t GetCount() const |
| 255 | { |
| 256 | return m_data.size(); |
| 257 | } |
| 258 | |
| 259 | // Gets a vector of all UNIQUE enum keys stored in the map. |
| 260 | std::vector<Enum> GetKeys() const |
| 261 | { |
| 262 | std::vector<Enum> keys; |
| 263 | Enum lastKey = static_cast<Enum>(-1); |
| 264 | bool first = true; |
| 265 | |
| 266 | for (const auto& pair : m_data) |
| 267 | { |
| 268 | if (first || pair.first != lastKey) |
| 269 | { |
| 270 | keys.push_back(pair.first); |
| 271 | lastKey = pair.first; |
| 272 | first = false; |
| 273 | } |
| 274 | } |
| 275 | |
| 276 | return keys; |
| 277 | } |
| 278 | |
| 279 | private: |
| 280 | // Helper to implement runtime type checking. |
| 281 | template <typename V, size_t... I> |
| 282 | bool IsMatchingTypeImpl(Enum e, std::index_sequence<I...>) const |
| 283 | { |
| 284 | bool result = false; |
| 285 | ((static_cast<size_t>(e) == I ? (result = std::is_same_v<std::remove_cvref_t<V>, mapping_t<static_cast<Enum>(I)>>, true) : false) || ...); |
| 286 | return result; |
| 287 | } |
| 288 | |
| 289 | // Helper to emplace at runtime-determined index |
| 290 | template <typename V, size_t... I> |
| 291 | void EmplaceAtRuntimeIndex(typename Variant::variant_t& variant, Enum e, V&& v, std::index_sequence<I...>) |
| 292 | { |
| 293 | size_t index = static_cast<size_t>(e) + 1; |
| 294 | bool handled = false; |
| 295 | |
| 296 | ( |
| 297 | [&] { |
| 298 | if (index == I + 1 && !handled) |
| 299 | { |
| 300 | using Emplacer = wsl::windows::wslc::EnumBasedVariantMapEmplacer<Enum, Mapping, V>; |
| 301 | Emplacer::template Emplace<I + 1>(variant, std::forward<V>(v)); |
| 302 | handled = true; |
| 303 | } |
| 304 | }(), |
| 305 | ...); |
| 306 | |
| 307 | if (!handled) |
| 308 | { |
| 309 | using CleanV = std::remove_cvref_t<V>; |
| 310 | THROW_HR_MSG(E_INVALIDARG, "Invalid enum value: %d", static_cast<int>(e)); |
| 311 | } |
| 312 | } |
| 313 | |
| 314 | std::multimap<Enum, typename Variant::variant_t> m_data; |
| 315 | }; |
| 316 | |
| 317 | // Helper for runtime emplacement into std::variant for EnumBasedVariantMap |
| 318 | template <typename Enum, template <Enum> typename Mapping, typename V> |
| 319 | struct EnumBasedVariantMapEmplacer |
| 320 | { |
| 321 | template <size_t Index> |
| 322 | static void Emplace(typename EnumBasedVariant<Enum, Mapping>::variant_t& variant, V&& value) |
| 323 | { |
| 324 | using TargetType = typename Mapping<static_cast<Enum>(Index - 1)>::value_t; |
| 325 | using CleanV = std::remove_cvref_t<V>; |
| 326 | |
| 327 | constexpr bool is_same_type = std::is_same_v<CleanV, TargetType>; |
| 328 | constexpr bool is_convertible = std::is_convertible_v<CleanV, TargetType>; |
| 329 | constexpr bool is_constructible = std::is_constructible_v<TargetType, CleanV>; |
| 330 | |
| 331 | if constexpr (is_same_type || is_convertible || is_constructible) |
| 332 | { |
| 333 | variant.template emplace<Index>(std::forward<V>(value)); |
| 334 | } |
| 335 | else |
| 336 | { |
| 337 | throw std::runtime_error("Runtime type mismatch: cannot convert value to target type for this enum value"); |
| 338 | } |
| 339 | } |
| 340 | }; |
| 341 | } // namespace wsl::windows::wslc |