1 | //===- IRMapping.h ----------------------------------------------*- C++ -*-===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This file defines a utility class for maintaining a mapping of SSA values, |
10 | // blocks, and operations. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #ifndef MLIR_IR_IRMAPPING_H |
15 | #define MLIR_IR_IRMAPPING_H |
16 | |
17 | #include "mlir/IR/Block.h" |
18 | |
19 | namespace mlir { |
20 | /// This is a utility class for mapping one set of IR entities to another. New |
21 | /// mappings can be inserted via 'map'. Existing mappings can be |
22 | /// found via the 'lookup*' functions. There are three variants that differ only |
23 | /// in return value when an existing is not found for the provided key: SSA |
24 | /// values, blocks, and operations. 'lookupOrNull' returns nullptr where as |
25 | /// 'lookupOrDefault' will return the lookup key. |
26 | class IRMapping { |
27 | public: |
28 | /// Inserts a new mapping for 'from' to 'to'. If there is an existing mapping, |
29 | /// it is overwritten. |
30 | void map(Value from, Value to) { valueMap[from] = to; } |
31 | void map(Block *from, Block *to) { blockMap[from] = to; } |
32 | void map(Operation *from, Operation *to) { operationMap[from] = to; } |
33 | |
34 | template <typename S, typename T, |
35 | std::enable_if_t<!std::is_assignable_v<Value, S> && |
36 | !std::is_assignable_v<Block *, S> && |
37 | !std::is_assignable_v<Operation *, S>> * = nullptr> |
38 | void map(S &&from, T &&to) { |
39 | for (auto [fromValue, toValue] : llvm::zip(from, to)) |
40 | map(fromValue, toValue); |
41 | } |
42 | |
43 | /// Erases a mapping for 'from'. |
44 | template <typename T> |
45 | void erase(T from) { |
46 | getMap<T>().erase(from); |
47 | } |
48 | |
49 | /// Checks to see if a mapping for 'from' exists. |
50 | template <typename T> |
51 | bool contains(T from) const { |
52 | return getMap<T>().count(from); |
53 | } |
54 | |
55 | /// Lookup a mapped value within the map. If a mapping for the provided value |
56 | /// does not exist then return nullptr. |
57 | template <typename T> |
58 | auto lookupOrNull(T from) const { |
59 | return lookupOrValue(from, T(nullptr)); |
60 | } |
61 | |
62 | /// Lookup a mapped value within the map. If a mapping for the provided value |
63 | /// does not exist then return the provided value. |
64 | template <typename T> |
65 | auto lookupOrDefault(T from) const { |
66 | return lookupOrValue(from, from); |
67 | } |
68 | |
69 | /// Lookup a mapped value within the map. This asserts the provided value |
70 | /// exists within the map. |
71 | template <typename T> |
72 | auto lookup(T from) const { |
73 | auto result = lookupOrNull(from); |
74 | assert(result && "expected 'from' to be contained within the map" ); |
75 | return result; |
76 | } |
77 | |
78 | /// Clears all mappings held by the mapper. |
79 | void clear() { valueMap.clear(); } |
80 | |
81 | /// Return the held value mapping. |
82 | const DenseMap<Value, Value> &getValueMap() const { return valueMap; } |
83 | |
84 | /// Return the held block mapping. |
85 | const DenseMap<Block *, Block *> &getBlockMap() const { return blockMap; } |
86 | |
87 | /// Return the held operation mapping. |
88 | const DenseMap<Operation *, Operation *> &getOperationMap() const { |
89 | return operationMap; |
90 | } |
91 | |
92 | private: |
93 | /// Return the map for the given value type. |
94 | template <typename T> |
95 | auto &getMap() const { |
96 | if constexpr (std::is_convertible_v<T, Value>) |
97 | return const_cast<DenseMap<Value, Value> &>(valueMap); |
98 | else if constexpr (std::is_convertible_v<T, Block *>) |
99 | return const_cast<DenseMap<Block *, Block *> &>(blockMap); |
100 | else |
101 | return const_cast<DenseMap<Operation *, Operation *> &>(operationMap); |
102 | } |
103 | |
104 | /// Utility lookupOrValue that looks up an existing key or returns the |
105 | /// provided value. |
106 | template <typename T> |
107 | auto lookupOrValue(T from, T value) const { |
108 | auto &map = getMap<T>(); |
109 | auto it = map.find(from); |
110 | return it != map.end() ? it->second : value; |
111 | } |
112 | |
113 | DenseMap<Value, Value> valueMap; |
114 | DenseMap<Block *, Block *> blockMap; |
115 | DenseMap<Operation *, Operation *> operationMap; |
116 | }; |
117 | |
118 | } // namespace mlir |
119 | |
120 | #endif // MLIR_IR_IRMAPPING_H |
121 | |