#pragma once

#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include "src/Accelerators/PIM/Common/PimCommon.hpp"

namespace onnx_mlir {

/**
 * \brief Get the offset of the ExtractSliceOp based on its static offsets and
 * its static tensor input.
 *
 * The static offsets represent the starting position of the slice in each
 * dimension, while the static tensor input gives its dimension size.
 *
 * \param sliceOp The ExtractSliceOp for which the actual offset needs to be
 * calculated.
 * \param inputShape The ShapedType of the ExtractSliceOp's input tensor
 * \return The actual offset of the ExtractSliceOp.
 */
size_t getSliceActualOffset(mlir::tensor::ExtractSliceOp& sliceOp, mlir::ShapedType& inputShape);

mlir::IntegerAttr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Value value);

template <class T>
size_t rangeLength(const mlir::iterator_range<T> range) {
  return std::distance(range.begin(), range.end());
}

/**
 * Retrieves the earliest operation that uses the given value within the value's
 * block.
 *
 * @param value The value for which to find the earliest user operation.
 * @return The earliest user operation that uses the given value within the
 * current block.
 */
mlir::Operation* getEarliestUserWithinBlock(mlir::Value value);

mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operation);

mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);

inline mlir::tensor::EmptyOp
createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {
  return mlir::tensor::EmptyOp::create(rewriter, loc, shapedType.getShape(), shapedType.getElementType());
}

} // namespace onnx_mlir