};

func::FuncOp GetMainFuncOp(ModuleOp module_op) {
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (func_op.getSymName() == "main") {
      return func_op;
    }
  }
  return {};
}

TEST_F(ConvertAssetArgsTest, ConvertsSingleAssetArg) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module {

  // MLIR @main function corresponds to the TF function "main_original".
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
      return %arg : tensor<1x2xf32>
    }
  )mlir");
  ASSERT_TRUE(module_op);

  absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases;

  constexpr StringRef kModuleOpStr =
      R"mlir(module attributes {tf_saved_model.semantics} {})mlir";

  Block block;
  ModuleOp module_op = ParseModuleOp(kModuleOpStr, block, ctx_);

  func::FuncOp init_func_op = GetInitializerFunction(
      module_op, /*initializer_type=*/kTfSavedModelInitializerInitType);

  EXPECT_THAT(init_func_op, IsNull());
}

TEST_F(TfSavedModelTest,

    module {
      func.func @main() -> () {
        return
      }
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleWithMainFunc);
  EXPECT_THAT(*module_op, NotNull());

  func::FuncOp main_func_op = FindMainFuncOp(*module_op);
  EXPECT_THAT(main_func_op, NotNull());
}

TEST_F(FindMainFuncOpTest, ReturnsNullWhenMainFuncOpIsPrivate) {

  // Runs calibration on `module_op` and returns a calibrated ModuleOp with
  // calibrated statistics embedded.
  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  // Exports `module_op` to SavedModel at `dst_saved_model_path`. This is used
  // to export the pre-calibrated `module_op` to SavedModel so that the

// passed to the target function.
LogicalResult LiftHashTableOpsToArguments(ModuleOp module_op,
                                          func::FuncOp target_func) {
  if (!llvm::hasSingleElement(target_func)) return success();
  if (!UsedBySupportedOps(module_op, target_func)) return success();
  if (IsMainOrInitializerFunction(module_op, target_func)) return success();

  llvm::StringMap<int> shared_name_to_arg_idx;

}  // namespace

absl::Status RunCalibrationPasses(
    mlir::ModuleOp module_op, MLIRContext& ctx,
    absl::string_view calibration_data_dir,
    const bool force_regenerate_calibration_data) {
  // Disable DumpTensor ops when running calibration.
  DisableDebugging(module_op);

  std::vector<std::string> skipping_aggregator_ops;
  if (!force_regenerate_calibration_data) {

SessionInitializerOp GetSessionInitializerOp(ModuleOp module_op);

// Returns the exported name for the session initializer function.
SmallVector<StringRef, 2> GetSessionInitializerExportedName(ModuleOp module_op);

// Returns initializer function ops. These functions' symbols are in the
// "initializers" attribute of the session initializer op.
SmallVector<func::FuncOp, 2> GetInitializerFunctions(ModuleOp module_op);

// `module_op` is the input graph to be quantized and it should contain
// StableHLO ops.
//
// Returns a quantized `ModuleOp` in StableHLO, potentially wrapped inside a
// XlaCallModuleOp. Returns a non-OK status if quantization fails, or any of
// `saved_model_bundle` or `quantization_py_function_lib` is a nullptr.
absl::StatusOr<mlir::ModuleOp> RunQuantization(

  mlir::OpBuilder builder(&ctx);
  auto module_op = builder.create<mlir::ModuleOp>(builder.getUnknownLoc());
  // Destroy by calling destroy() to avoid memory leak since it is allocated
  // with malloc().
  const absl::Cleanup module_op_cleanup = [module_op] { module_op->destroy(); };

  const mlir::LogicalResult result = pm.run(module_op);
  EXPECT_FALSE(failed(result));