#pragma once

#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Builders.h"

#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/SmallVector.h"

namespace onnx_mlir {

struct NormalizedLoopResult {
  mlir::Value inductionVar;
  llvm::SmallVector<mlir::Value, 4> results;
  mlir::scf::ForOp loop;

  bool wasInlined() const { return !loop; }
};

using NormalizedLoopBodyBuilder = llvm::function_ref<mlir::LogicalResult(
  mlir::OpBuilder&, mlir::Location, mlir::Value, mlir::ValueRange, llvm::SmallVectorImpl<mlir::Value>&)>;

mlir::FailureOr<NormalizedLoopResult> buildNormalizedScfFor(mlir::OpBuilder& builder,
                                                            mlir::Location loc,
                                                            mlir::Value lowerBound,
                                                            mlir::Value upperBound,
                                                            mlir::Value step,
                                                            mlir::ValueRange initArgs,
                                                            NormalizedLoopBodyBuilder bodyBuilder);

} // namespace onnx_mlir