// Legalize TF While to TFL While with calls to the original functions from the
// cond and body regions.
struct LegalizeWhilePass
    : public impl::LegalizeWhilePassBase<LegalizeWhilePass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeWhilePass)
  void RunOnFunction(func::FuncOp func);

  void runOnOperation() override {
    for (auto op : getOperation().getOps<func::FuncOp>()) RunOnFunction(op);

      Option<"preserve_assert_op_", "preserve-assert-op",
             "bool", "false", "Preserve AssertOp during tfl legalization.">,
  ];
}

def LegalizeWhilePass : Pass<"tfl-legalize-tf-while", "mlir::ModuleOp"> {
  let summary = "Legalize from TensorFlow While to TensorFlow Lite While.";
  let constructor = "CreateLegalizeTFWhilePass()";