SymbolTable.h

Go to the documentation of this file.

//===- SymbolTable.h - MLIR Symbol Table Class ------------------*- 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_SYMBOLTABLE_H
#define MLIR_IR_SYMBOLTABLE_H

#include "mlir/IR/OpDefinition.h"
#include "llvm/ADT/StringMap.h"

namespace mlir {
class Identifier;
class Operation;

class SymbolTable {
public:
  SymbolTable(Operation *symbolTableOp);

  Operation *lookup(StringRef name) const;
  template <typename T> T lookup(StringRef name) const {
    return dyn_cast_or_null<T>(lookup(name));
  }

  void erase(Operation *symbol);

  void insert(Operation *symbol, Block::iterator insertPt = {});

  static StringRef getSymbolAttrName() { return "sym_name"; }

  Operation *getOp() const { return symbolTableOp; }

  //===--------------------------------------------------------------------===//
  // Symbol Utilities
  //===--------------------------------------------------------------------===//

  static Operation *lookupSymbolIn(Operation *op, StringRef symbol);

  static Operation *lookupNearestSymbolFrom(Operation *from, StringRef symbol);

  class SymbolUse {
  public:
    SymbolUse(Operation *op, SymbolRefAttr symbolRef)
        : owner(op), symbolRef(symbolRef) {}

    Operation *getUser() const { return owner; }

    SymbolRefAttr getSymbolRef() const { return symbolRef; }

  private:
    Operation *owner;

    SymbolRefAttr symbolRef;
  };

  class UseRange {
  public:
    UseRange(std::vector<SymbolUse> &&uses) : uses(std::move(uses)) {}

    using iterator = std::vector<SymbolUse>::const_iterator;
    iterator begin() const { return uses.begin(); }
    iterator end() const { return uses.end(); }

  private:
    std::vector<SymbolUse> uses;
  };

  static Optional<UseRange> getSymbolUses(Operation *from);

  static Optional<UseRange> getSymbolUses(StringRef symbol, Operation *from);

  static bool symbolKnownUseEmpty(StringRef symbol, Operation *from);

  LLVM_NODISCARD static LogicalResult replaceAllSymbolUses(StringRef oldSymbol,
                                                           StringRef newSymbol,
                                                           Operation *from);

private:
  Operation *symbolTableOp;

  llvm::StringMap<Operation *> symbolTable;

  unsigned uniquingCounter = 0;
};

//===----------------------------------------------------------------------===//
// SymbolTable Trait Types
//===----------------------------------------------------------------------===//

namespace OpTrait {
namespace impl {
LogicalResult verifySymbolTable(Operation *op);
LogicalResult verifySymbol(Operation *op);
} // namespace impl

template <typename ConcreteType>
class SymbolTable : public TraitBase<ConcreteType, SymbolTable> {
public:
  static LogicalResult verifyTrait(Operation *op) {
    return impl::verifySymbolTable(op);
  }

  Operation *lookupSymbol(StringRef name) {
    return mlir::SymbolTable::lookupSymbolIn(this->getOperation(), name);
  }
  template <typename T> T lookupSymbol(StringRef name) {
    return dyn_cast_or_null<T>(lookupSymbol(name));
  }
};

template <typename ConcreteType>
class Symbol : public TraitBase<ConcreteType, Symbol> {
public:
  static LogicalResult verifyTrait(Operation *op) {
    return impl::verifySymbol(op);
  }

  StringRef getName() {
    return this->getOperation()
        ->template getAttrOfType<StringAttr>(
            mlir::SymbolTable::getSymbolAttrName())
        .getValue();
  }

  void setName(StringRef name) {
    this->getOperation()->setAttr(
        mlir::SymbolTable::getSymbolAttrName(),
        StringAttr::get(name, this->getOperation()->getContext()));
  }

  Optional<::mlir::SymbolTable::UseRange> getSymbolUses(Operation *from) {
    return ::mlir::SymbolTable::getSymbolUses(getName(), from);
  }

  bool symbolKnownUseEmpty(Operation *from) {
    return ::mlir::SymbolTable::symbolKnownUseEmpty(getName(), from);
  }

  LLVM_NODISCARD LogicalResult replaceAllSymbolUses(StringRef newSymbol,
                                                    Operation *from) {
    return ::mlir::SymbolTable::replaceAllSymbolUses(getName(), newSymbol,
                                                     from);
  }
};

} // end namespace OpTrait
} // end namespace mlir

#endif // MLIR_IR_SYMBOLTABLE_H