#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

#include "ShapingUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {

bool isShapingOnlyOp(Operation *op) {
  return isa<tensor::CastOp,
             tensor::CollapseShapeOp,
             tensor::ExpandShapeOp,
             tensor::ExtractSliceOp,
             tensor::InsertSliceOp,
             tensor::ConcatOp,
             tensor::EmptyOp,
             tensor::ExtractOp,
             tensor::InsertOp,
             tensor::SplatOp,
             linalg::TransposeOp,
             ONNXTransposeOp,
             spatial::SpatConcatOp,
             spatial::SpatExtractRowsOp>(op);
}

bool isPureIndexComputationOp(Operation *op) {
  if (op->getNumRegions() != 0 || op->getNumResults() == 0 || op->hasTrait<OpTrait::IsTerminator>()
      || !isMemoryEffectFree(op))
    return false;
  auto isIndexOrInteger = [](Type type) { return type.isIndex() || isa<IntegerType>(type); };
  return llvm::all_of(op->getOperandTypes(), isIndexOrInteger)
      && llvm::all_of(op->getResultTypes(), isIndexOrInteger);
}

} // namespace onnx_mlir