1 //===- BlockAndValueMapping.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 for multiple
10 // value types.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_IR_BLOCKANDVALUEMAPPING_H
15 #define MLIR_IR_BLOCKANDVALUEMAPPING_H
16 
17 #include "mlir/IR/Block.h"
18 
19 namespace mlir {
20 // This is a utility class for mapping one set of values to another. New
21 // mappings can be inserted via 'map'. Existing mappings can be
22 // found via the 'lookup*' functions. There are two variants that differ only in
23 // return value when an existing is not found for the provided key.
24 // 'lookupOrNull' returns nullptr where as 'lookupOrDefault' will return the
25 // lookup key.
26 class BlockAndValueMapping {
27 public:
28   /// Inserts a new mapping for 'from' to 'to'. If there is an existing mapping,
29   /// it is overwritten.
map(Block * from,Block * to)30   void map(Block *from, Block *to) { blockMap[from] = to; }
map(Value from,Value to)31   void map(Value from, Value to) { valueMap[from] = to; }
32 
33   template <
34       typename S, typename T,
35       std::enable_if_t<!std::is_assignable<Value, S>::value &&
36                        !std::is_assignable<Block *, S>::value> * = nullptr>
map(S && from,T && to)37   void map(S &&from, T &&to) {
38     for (auto pair : llvm::zip(from, to))
39       map(std::get<0>(pair), std::get<1>(pair));
40   }
41 
42   /// Erases a mapping for 'from'.
erase(Block * from)43   void erase(Block *from) { blockMap.erase(from); }
erase(Value from)44   void erase(Value from) { valueMap.erase(from); }
45 
46   /// Checks to see if a mapping for 'from' exists.
contains(Block * from)47   bool contains(Block *from) const { return blockMap.count(from); }
contains(Value from)48   bool contains(Value from) const { return valueMap.count(from); }
49 
50   /// Lookup a mapped value within the map. If a mapping for the provided value
51   /// does not exist then return nullptr.
lookupOrNull(Block * from)52   Block *lookupOrNull(Block *from) const {
53     return lookupOrValue(from, (Block *)nullptr);
54   }
lookupOrNull(Value from)55   Value lookupOrNull(Value from) const { return lookupOrValue(from, Value()); }
56 
57   /// Lookup a mapped value within the map. If a mapping for the provided value
58   /// does not exist then return the provided value.
lookupOrDefault(Block * from)59   Block *lookupOrDefault(Block *from) const {
60     return lookupOrValue(from, from);
61   }
lookupOrDefault(Value from)62   Value lookupOrDefault(Value from) const { return lookupOrValue(from, from); }
63 
64   /// Lookup a mapped value within the map. This asserts the provided value
65   /// exists within the map.
66   template <typename T>
lookup(T from)67   T lookup(T from) const {
68     auto result = lookupOrNull(from);
69     assert(result && "expected 'from' to be contained within the map");
70     return result;
71   }
72 
73   /// Clears all mappings held by the mapper.
clear()74   void clear() { valueMap.clear(); }
75 
76   /// Return the held value mapping.
getValueMap()77   const DenseMap<Value, Value> &getValueMap() const { return valueMap; }
78 
79   /// Return the held block mapping.
getBlockMap()80   const DenseMap<Block *, Block *> &getBlockMap() const { return blockMap; }
81 
82 private:
83   /// Utility lookupOrValue that looks up an existing key or returns the
84   /// provided value.
lookupOrValue(Block * from,Block * value)85   Block *lookupOrValue(Block *from, Block *value) const {
86     auto it = blockMap.find(from);
87     return it != blockMap.end() ? it->second : value;
88   }
lookupOrValue(Value from,Value value)89   Value lookupOrValue(Value from, Value value) const {
90     auto it = valueMap.find(from);
91     return it != valueMap.end() ? it->second : value;
92   }
93 
94   DenseMap<Value, Value> valueMap;
95   DenseMap<Block *, Block *> blockMap;
96 };
97 
98 } // namespace mlir
99 
100 #endif // MLIR_IR_BLOCKANDVALUEMAPPING_H
101