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
19namespace 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.
26class IRMapping {
27public:
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
92private:
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

source code of mlir/include/mlir/IR/IRMapping.h