DialectInterface.h

Go to the documentation of this file.

//===- DialectInterface.h - IR Dialect Interfaces ---------------*- 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
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_IR_DIALECTINTERFACE_H
#define MLIR_IR_DIALECTINTERFACE_H

#include "mlir/Support/STLExtras.h"
#include "llvm/ADT/DenseSet.h"

namespace mlir {
class Dialect;
class MLIRContext;
class Operation;

//===----------------------------------------------------------------------===//
// DialectInterface
//===----------------------------------------------------------------------===//
namespace detail {
template <typename ConcreteType, typename BaseT>
class DialectInterfaceBase : public BaseT {
public:
  using Base = DialectInterfaceBase<ConcreteType, BaseT>;

  static ClassID *getInterfaceID() { return ClassID::getID<ConcreteType>(); }

protected:
  DialectInterfaceBase(Dialect *dialect) : BaseT(dialect, getInterfaceID()) {}
};
} // end namespace detail

class DialectInterface {
public:
  virtual ~DialectInterface();

  template <typename ConcreteType>
  using Base = detail::DialectInterfaceBase<ConcreteType, DialectInterface>;

  Dialect *getDialect() const { return dialect; }

  ClassID *getID() const { return interfaceID; }

protected:
  DialectInterface(Dialect *dialect, ClassID *id)
      : dialect(dialect), interfaceID(id) {}

private:
  Dialect *dialect;

  ClassID *interfaceID;
};

//===----------------------------------------------------------------------===//
// DialectInterfaceCollection
//===----------------------------------------------------------------------===//

namespace detail {
class DialectInterfaceCollectionBase {
  struct InterfaceKeyInfo : public DenseMapInfo<const DialectInterface *> {
    using DenseMapInfo<const DialectInterface *>::isEqual;

    static unsigned getHashValue(Dialect *key) { return llvm::hash_value(key); }
    static unsigned getHashValue(const DialectInterface *key) {
      return getHashValue(key->getDialect());
    }

    static bool isEqual(Dialect *lhs, const DialectInterface *rhs) {
      if (rhs == getEmptyKey() || rhs == getTombstoneKey())
        return false;
      return lhs == rhs->getDialect();
    }
  };

  using InterfaceSetT = DenseSet<const DialectInterface *, InterfaceKeyInfo>;
  using InterfaceVectorT = std::vector<const DialectInterface *>;

public:
  DialectInterfaceCollectionBase(MLIRContext *ctx, ClassID *interfaceKind);
  virtual ~DialectInterfaceCollectionBase();

protected:
  const DialectInterface *getInterfaceFor(Operation *op) const;

  const DialectInterface *getInterfaceFor(Dialect *dialect) const {
    auto it = interfaces.find_as(dialect);
    return it == interfaces.end() ? nullptr : *it;
  }

  template <typename InterfaceT>
  class iterator : public llvm::mapped_iterator<
                       InterfaceVectorT::const_iterator,
                       const InterfaceT &(*)(const DialectInterface *)> {
    static const InterfaceT &remapIt(const DialectInterface *interface) {
      return *static_cast<const InterfaceT *>(interface);
    }

    iterator(InterfaceVectorT::const_iterator it)
        : llvm::mapped_iterator<
              InterfaceVectorT::const_iterator,
              const InterfaceT &(*)(const DialectInterface *)>(it, &remapIt) {}

    friend DialectInterfaceCollectionBase;
  };

  template <typename InterfaceT> iterator<InterfaceT> interface_begin() const {
    return iterator<InterfaceT>(orderedInterfaces.begin());
  }
  template <typename InterfaceT> iterator<InterfaceT> interface_end() const {
    return iterator<InterfaceT>(orderedInterfaces.end());
  }

private:
  InterfaceSetT interfaces;
  // NOTE: SetVector does not provide find access, so it can't be used here.
  InterfaceVectorT orderedInterfaces;
};
} // namespace detail

template <typename InterfaceType>
class DialectInterfaceCollection
    : public detail::DialectInterfaceCollectionBase {
public:
  using Base = DialectInterfaceCollection<InterfaceType>;

  DialectInterfaceCollection(MLIRContext *ctx)
      : detail::DialectInterfaceCollectionBase(
            ctx, InterfaceType::getInterfaceID()) {}

  template <typename Object>
  const InterfaceType *getInterfaceFor(Object *obj) const {
    return static_cast<const InterfaceType *>(
        detail::DialectInterfaceCollectionBase::getInterfaceFor(obj));
  }

  using iterator =
      detail::DialectInterfaceCollectionBase::iterator<InterfaceType>;
  iterator begin() const { return interface_begin<InterfaceType>(); }
  iterator end() const { return interface_end<InterfaceType>(); }

private:
  using detail::DialectInterfaceCollectionBase::interface_begin;
  using detail::DialectInterfaceCollectionBase::interface_end;
};

} // namespace mlir

#endif