#pragma once

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
#include "mlir/Transforms/FoldUtils.h"

#include "llvm/ADT/DenseMap.h"

#include <optional>

namespace onnx_mlir {

class ConstantPool {
public:
  ConstantPool(mlir::Operation *constantAnchor, mlir::OpBuilder &builder);

  mlir::Value getIndex(int64_t value);
  mlir::Value get(mlir::Type type, mlir::Attribute value);

private:
  mlir::Operation *anchor;
  mlir::Block *block;
  mlir::OpBuilder &builder;
  llvm::DenseMap<std::pair<mlir::Type, mlir::Attribute>, mlir::Value> cache;
};

mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp);

mlir::Value
getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);

mlir::Value
getOrCreateConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);

mlir::Value
createConstantAtHostBlockStart(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::TypedAttr value);

mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);

mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value);

mlir::Value getOrCreateIndexConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, int64_t value);

void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter);

std::optional<int64_t> matchConstantIndexValue(mlir::Value value);

std::optional<int64_t> matchConstantIndexValue(mlir::OpFoldResult value);

} // namespace onnx_mlir