121 lines
4.1 KiB
C++
121 lines
4.1 KiB
C++
//===- IRMapping.h ----------------------------------------------*- C++ -*-===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file defines a utility class for maintaining a mapping of SSA values,
|
|
// blocks, and operations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_IR_IRMAPPING_H
|
|
#define MLIR_IR_IRMAPPING_H
|
|
|
|
#include "mlir/IR/Block.h"
|
|
|
|
namespace mlir {
|
|
/// This is a utility class for mapping one set of IR entities to another. New
|
|
/// mappings can be inserted via 'map'. Existing mappings can be
|
|
/// found via the 'lookup*' functions. There are three variants that differ only
|
|
/// in return value when an existing is not found for the provided key: SSA
|
|
/// values, blocks, and operations. 'lookupOrNull' returns nullptr where as
|
|
/// 'lookupOrDefault' will return the lookup key.
|
|
class IRMapping {
|
|
public:
|
|
/// Inserts a new mapping for 'from' to 'to'. If there is an existing mapping,
|
|
/// it is overwritten.
|
|
void map(Value from, Value to) { valueMap[from] = to; }
|
|
void map(Block *from, Block *to) { blockMap[from] = to; }
|
|
void map(Operation *from, Operation *to) { operationMap[from] = to; }
|
|
|
|
template <typename S, typename T,
|
|
std::enable_if_t<!std::is_assignable_v<Value, S> &&
|
|
!std::is_assignable_v<Block *, S> &&
|
|
!std::is_assignable_v<Operation *, S>> * = nullptr>
|
|
void map(S &&from, T &&to) {
|
|
for (auto [fromValue, toValue] : llvm::zip(from, to))
|
|
map(fromValue, toValue);
|
|
}
|
|
|
|
/// Erases a mapping for 'from'.
|
|
template <typename T>
|
|
void erase(T from) {
|
|
getMap<T>().erase(from);
|
|
}
|
|
|
|
/// Checks to see if a mapping for 'from' exists.
|
|
template <typename T>
|
|
bool contains(T from) const {
|
|
return getMap<T>().count(from);
|
|
}
|
|
|
|
/// Lookup a mapped value within the map. If a mapping for the provided value
|
|
/// does not exist then return nullptr.
|
|
template <typename T>
|
|
auto lookupOrNull(T from) const {
|
|
return lookupOrValue(from, T(nullptr));
|
|
}
|
|
|
|
/// Lookup a mapped value within the map. If a mapping for the provided value
|
|
/// does not exist then return the provided value.
|
|
template <typename T>
|
|
auto lookupOrDefault(T from) const {
|
|
return lookupOrValue(from, from);
|
|
}
|
|
|
|
/// Lookup a mapped value within the map. This asserts the provided value
|
|
/// exists within the map.
|
|
template <typename T>
|
|
auto lookup(T from) const {
|
|
auto result = lookupOrNull(from);
|
|
assert(result && "expected 'from' to be contained within the map");
|
|
return result;
|
|
}
|
|
|
|
/// Clears all mappings held by the mapper.
|
|
void clear() { valueMap.clear(); }
|
|
|
|
/// Return the held value mapping.
|
|
const DenseMap<Value, Value> &getValueMap() const { return valueMap; }
|
|
|
|
/// Return the held block mapping.
|
|
const DenseMap<Block *, Block *> &getBlockMap() const { return blockMap; }
|
|
|
|
/// Return the held operation mapping.
|
|
const DenseMap<Operation *, Operation *> &getOperationMap() const {
|
|
return operationMap;
|
|
}
|
|
|
|
private:
|
|
/// Return the map for the given value type.
|
|
template <typename T>
|
|
auto &getMap() const {
|
|
if constexpr (std::is_convertible_v<T, Value>)
|
|
return const_cast<DenseMap<Value, Value> &>(valueMap);
|
|
else if constexpr (std::is_convertible_v<T, Block *>)
|
|
return const_cast<DenseMap<Block *, Block *> &>(blockMap);
|
|
else
|
|
return const_cast<DenseMap<Operation *, Operation *> &>(operationMap);
|
|
}
|
|
|
|
/// Utility lookupOrValue that looks up an existing key or returns the
|
|
/// provided value.
|
|
template <typename T>
|
|
auto lookupOrValue(T from, T value) const {
|
|
auto &map = getMap<T>();
|
|
auto it = map.find(from);
|
|
return it != map.end() ? it->second : value;
|
|
}
|
|
|
|
DenseMap<Value, Value> valueMap;
|
|
DenseMap<Block *, Block *> blockMap;
|
|
DenseMap<Operation *, Operation *> operationMap;
|
|
};
|
|
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_IR_IRMAPPING_H
|