generator_region.getNextFuncOtherArgs(),
             /*use_region_args=*/true, /*forward_block_args=*/false);
  CreateCall(generator_op, finalize_function, generator_region.getFinalize(),
             generator_region.getFinalizeFuncOtherArgs(),
             /*use_region_args=*/true, /*forward_block_args=*/false);

  generator_op->replaceAllUsesWith(generator_region->getResults());

                                       /*allow_return_of_existing=*/true);

  finalize_name = GetName(regional, "_finalize");
  finalize_name =
      ExtractSingleBlockRegion(symbol_table, regional.getFinalize(),
                               finalize_name, extern_values, worklist,
                               /*extern_values_passthrough=*/false,
                               /*only_one_return_value=*/false,

    // and `finalize`.
    regions.push_back(RegionSuccessor(
        &getNext(), getNext().front().getArguments().slice(0, num)));
    regions.push_back(RegionSuccessor(
        &getFinalize(), getFinalize().front().getArguments().slice(0, num)));
  } else {
    // `finalize` branches back to the op itself, not passing any arguments.
    regions.push_back(RegionSuccessor());
  }
}

          this->getOperation()->getOperands().size() - 1);
    }
  } else if (auto regionOp = llvm::dyn_cast<GeneratorDatasetRegionOp>(
                 this->getOperation()->getParentOp())) {
    if (&regionOp.getFinalize() == this->getOperation()->getParentRegion()) {
      // `finalize`'s returns get discarded.
      return MutableOperandRange(this->getOperation(), 0, 0);
    }
  }