StorageUniquer.h

Go to the documentation of this file.

//===- StorageUniquer.h - Common Storage Class Uniquer ----------*- 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_SUPPORT_STORAGEUNIQUER_H
#define MLIR_SUPPORT_STORAGEUNIQUER_H

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

namespace mlir {
namespace detail {
struct StorageUniquerImpl;

template <typename ImplTy, typename... Args>
using has_impltype_getkey_t = decltype(ImplTy::getKey(std::declval<Args>()...));

template <typename ImplTy, typename T>
using has_impltype_hash_t = decltype(ImplTy::hashKey(std::declval<T>()));
} // namespace detail

class StorageUniquer {
public:
  StorageUniquer();
  ~StorageUniquer();

  class BaseStorage {
  public:
    unsigned getKind() const { return kind; }

  protected:
    BaseStorage() : kind(0) {}

  private:
    friend detail::StorageUniquerImpl;

    unsigned kind;
  };

  class StorageAllocator {
  public:
    template <typename T> ArrayRef<T> copyInto(ArrayRef<T> elements) {
      if (elements.empty())
        return llvm::None;
      auto result = allocator.Allocate<T>(elements.size());
      std::uninitialized_copy(elements.begin(), elements.end(), result);
      return ArrayRef<T>(result, elements.size());
    }

    StringRef copyInto(StringRef str) {
      auto result = copyInto(ArrayRef<char>(str.data(), str.size()));
      return StringRef(result.data(), str.size());
    }

    template <typename T> T *allocate() { return allocator.Allocate<T>(); }

    void *allocate(size_t size, size_t alignment) {
      return allocator.Allocate(size, alignment);
    }

  private:
    llvm::BumpPtrAllocator allocator;
  };

  template <typename Storage, typename Arg, typename... Args>
  Storage *get(std::function<void(Storage *)> initFn, unsigned kind, Arg &&arg,
               Args &&... args) {
    // Construct a value of the derived key type.
    auto derivedKey =
        getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);

    // Create a hash of the kind and the derived key.
    unsigned hashValue = getHash<Storage>(kind, derivedKey);

    // Generate an equality function for the derived storage.
    std::function<bool(const BaseStorage *)> isEqual =
        [&derivedKey](const BaseStorage *existing) {
          return static_cast<const Storage &>(*existing) == derivedKey;
        };

    // Generate a constructor function for the derived storage.
    std::function<BaseStorage *(StorageAllocator &)> ctorFn =
        [&](StorageAllocator &allocator) {
          auto *storage = Storage::construct(allocator, derivedKey);
          if (initFn)
            initFn(storage);
          return storage;
        };

    // Get an instance for the derived storage.
    return static_cast<Storage *>(getImpl(kind, hashValue, isEqual, ctorFn));
  }

  template <typename Storage>
  Storage *get(std::function<void(Storage *)> initFn, unsigned kind) {
    auto ctorFn = [&](StorageAllocator &allocator) {
      auto *storage = new (allocator.allocate<Storage>()) Storage();
      if (initFn)
        initFn(storage);
      return storage;
    };
    return static_cast<Storage *>(getImpl(kind, ctorFn));
  }

  template <typename Storage, typename Arg, typename... Args>
  void erase(unsigned kind, Arg &&arg, Args &&... args) {
    // Construct a value of the derived key type.
    auto derivedKey =
        getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);

    // Create a hash of the kind and the derived key.
    unsigned hashValue = getHash<Storage>(kind, derivedKey);

    // Generate an equality function for the derived storage.
    std::function<bool(const BaseStorage *)> isEqual =
        [&derivedKey](const BaseStorage *existing) {
          return static_cast<const Storage &>(*existing) == derivedKey;
        };

    // Attempt to erase the storage instance.
    eraseImpl(kind, hashValue, isEqual, [](BaseStorage *storage) {
      static_cast<Storage *>(storage)->cleanup();
    });
  }

private:
  BaseStorage *getImpl(unsigned kind, unsigned hashValue,
                       function_ref<bool(const BaseStorage *)> isEqual,
                       std::function<BaseStorage *(StorageAllocator &)> ctorFn);

  BaseStorage *getImpl(unsigned kind,
                       std::function<BaseStorage *(StorageAllocator &)> ctorFn);

  void eraseImpl(unsigned kind, unsigned hashValue,
                 function_ref<bool(const BaseStorage *)> isEqual,
                 std::function<void(BaseStorage *)> cleanupFn);

  std::unique_ptr<detail::StorageUniquerImpl> impl;

  //===--------------------------------------------------------------------===//
  // Key Construction
  //===--------------------------------------------------------------------===//

  template <typename ImplTy, typename... Args>
  static typename std::enable_if<
      is_detected<detail::has_impltype_getkey_t, ImplTy, Args...>::value,
      typename ImplTy::KeyTy>::type
  getKey(Args &&... args) {
    return ImplTy::getKey(args...);
  }
  template <typename ImplTy, typename... Args>
  static typename std::enable_if<
      !is_detected<detail::has_impltype_getkey_t, ImplTy, Args...>::value,
      typename ImplTy::KeyTy>::type
  getKey(Args &&... args) {
    return typename ImplTy::KeyTy(args...);
  }

  //===--------------------------------------------------------------------===//
  // Key and Kind Hashing
  //===--------------------------------------------------------------------===//

  template <typename ImplTy, typename DerivedKey>
  static typename std::enable_if<
      is_detected<detail::has_impltype_hash_t, ImplTy, DerivedKey>::value,
      ::llvm::hash_code>::type
  getHash(unsigned kind, const DerivedKey &derivedKey) {
    return llvm::hash_combine(kind, ImplTy::hashKey(derivedKey));
  }
  template <typename ImplTy, typename DerivedKey>
  static typename std::enable_if<
      !is_detected<detail::has_impltype_hash_t, ImplTy, DerivedKey>::value,
      ::llvm::hash_code>::type
  getHash(unsigned kind, const DerivedKey &derivedKey) {
    return llvm::hash_combine(
        kind, DenseMapInfo<DerivedKey>::getHashValue(derivedKey));
  }
};
} // end namespace mlir

#endif