)mlir";

TEST_F(AttrsAndConstraintsTest, HasStaticShapeSucceedsWithStaticShapes) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStatic);
  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  Value dot_general_result =
      FindOperationOfType<DotGeneralOp>(main_fn)->getResult(0);
  EXPECT_TRUE(HasStaticShape(dot_general_result));

        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) {
  constexpr absl::string_view kModuleWithPrivateMainFunc = R"mlir(
    module {

  }
)mlir";

TEST_F(LiftAsFunctionCallTest, FunctionLiftedAsXlaCallModuleOp) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStableHlo);
  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  auto dot_general_op =
      FindOperationOfType<mlir::stablehlo::DotGeneralOp>(main_fn);

  const SmallVector<NamedAttribute>& attributes = {

}

TEST_F(ApplyQuantizationParamsPropagationTest,
       ConstsUsedMultipleTimesAreDuplicated) {
  const OwningOpRef<ModuleOp> module_op_ref =
      ParseModuleOpString(kModuleTFLite);
  func::FuncOp main_fn = FindMainFuncOp(*module_op_ref);

  auto op_quant_spec_getter = [&](Operation* op) {
    return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
  };
  QuantizationDriver quantization_driver(

// `main` or `serving_default`. If `func::FuncOps` with both names exist, the
// function with name "main" takes precedence. Returns null if no such a
// function exists.
func::FuncOp FindMainFuncOp(ModuleOp module_op);

}  // namespace mlir::quant

bool IsPublicFuncOp(func::FuncOp func_op) {
  return SymbolTable::getSymbolVisibility(&*func_op) ==
         SymbolTable::Visibility::Public;
}

}  // namespace

func::FuncOp FindMainFuncOp(ModuleOp module_op) {
  if (const auto main_func_op = module_op.lookupSymbol<func::FuncOp>(
          kImportModelDefaultGraphFuncName);
      main_func_op != nullptr && IsPublicFuncOp(main_func_op)) {

  // Returns `failure` iff there is no `@main` or no such operation is found in
  // `@main`.
  template <typename OpT>
  FailureOr<OpT> FindFirstOpFromMainFunc(ModuleOp module_op) {
    func::FuncOp main_func_op = FindMainFuncOp(module_op);
    if (main_func_op == nullptr) return failure();

    auto ops = main_func_op.getOps<OpT>();
    if (ops.empty()) return failure();

    return *ops.begin();
  }

  main_func_op.setArgAttrs(arg_idx, new_arg_attrs);
}

}  // namespace

FailureOr<SmallVector<AssetFileDef>> ConvertAssetArgs(ModuleOp module_op) {
  func::FuncOp main_func_op = FindMainFuncOp(module_op);
  if (!main_func_op) return failure();

  SmallVector<StringRef> input_names = GetEntryFunctionInputs(main_func_op);
  SymbolTable symbol_table(module_op);
  SmallVector<AssetFileDef> asset_file_defs;

        func.return %2 : tensor<2x3x2x2xf32>
      }
    }
  )mlir";
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleGather);
  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  Operation* gather_op = FindOperationOfType<GatherOp>(main_fn);
  const auto spec = GetStableHloQuantConstraints(gather_op);

  EXPECT_THAT(spec, NotNull());

}

void MergeInitializerFunctionOpsToMainPass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext* ctx = module_op.getContext();

  func::FuncOp main_func_op = FindMainFuncOp(module_op);
  if (!main_func_op) {
    module_op.emitError("Main function op not found.");
    return signalPassFailure();
  }

  GraphOp main_graph_op = GetGraphOpFromFuncOp(main_func_op);