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. map(Value from,Value to)30 void map(Value from, Value to) { valueMap[from] = to; } map(Block * from,Block * to)31 void map(Block *from, Block *to) { blockMap[from] = to; } map(Operation * from,Operation * 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> map(S && from,T && to)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> erase(T from)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> contains(T from)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> lookupOrNull(T from)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> lookupOrDefault(T from)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> lookup(T from)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. clear()79 void clear() { valueMap.clear(); } 80 81 /// Return the held value mapping. getValueMap()82 const DenseMap<Value, Value> &getValueMap() const { return valueMap; } 83 84 /// Return the held block mapping. getBlockMap()85 const DenseMap<Block *, Block *> &getBlockMap() const { return blockMap; } 86 87 /// Return the held operation mapping. getOperationMap()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> getMap()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> lookupOrValue(T from,T value)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