BlockAndValueMapping.h

Go to the documentation of this file.

//===- BlockAndValueMapping.h -----------------------------------*- C++ -*-===//
//
// Part of the MLIR 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 for multiple
// value types.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_IR_BLOCKANDVALUEMAPPING_H
#define MLIR_IR_BLOCKANDVALUEMAPPING_H

#include "mlir/IR/Block.h"

namespace mlir {
// This is a utility class for mapping one set of values to another. New
// mappings can be inserted via 'map'. Existing mappings can be
// found via the 'lookup*' functions. There are two variants that differ only in
// return value when an existing is not found for the provided key.
// 'lookupOrNull' returns nullptr where as 'lookupOrDefault' will return the
// lookup key.
class BlockAndValueMapping {
public:
  void map(Block *from, Block *to) { valueMap[from] = to; }
  void map(Value from, Value to) {
    valueMap[from.getAsOpaquePointer()] = to.getAsOpaquePointer();
  }

  void erase(Block *from) { valueMap.erase(from); }
  void erase(Value from) { valueMap.erase(from.getAsOpaquePointer()); }

  bool contains(Block *from) const { return valueMap.count(from); }
  bool contains(Value from) const {
    return valueMap.count(from.getAsOpaquePointer());
  }

  Block *lookupOrNull(Block *from) const {
    return lookupOrValue(from, (Block *)nullptr);
  }
  Value lookupOrNull(Value from) const { return lookupOrValue(from, Value()); }

  Block *lookupOrDefault(Block *from) const {
    return lookupOrValue(from, from);
  }
  Value lookupOrDefault(Value from) const { return lookupOrValue(from, from); }

  template <typename T> T lookup(T from) const {
    auto result = lookupOrNull(from);
    assert(result && "expected 'from' to be contained within the map");
    return result;
  }

  void clear() { valueMap.clear(); }

private:
  Block *lookupOrValue(Block *from, Block *value) const {
    auto it = valueMap.find(from);
    return it != valueMap.end() ? reinterpret_cast<Block *>(it->second) : value;
  }
  Value lookupOrValue(Value from, Value value) const {
    auto it = valueMap.find(from.getAsOpaquePointer());
    return it != valueMap.end() ? Value::getFromOpaquePointer(it->second)
                                : value;
  }

  DenseMap<void *, void *> valueMap;
};

} // end namespace mlir

#endif // MLIR_IR_BLOCKANDVALUEMAPPING_H