#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace mlir::quant {
namespace {

using ::mlir::tf_saved_model::AssetOp;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::LookupBoundInputOfType;
using ::tensorflow::AssetFileDef;

// Given argument attributes `arg_attrs`, returns a new set of argument
// attributes where the "tf_saved_model.bound_input" attribute has been replaced

  llvm::DenseSet<Value> seen;
  BlockArgument arg;
  if (traceUpwardsToArgument(resource, seen, &arg).failed()) return nullptr;
  Operation *global =
      tf_saved_model::LookupBoundInputOfType<tf_saved_model::GlobalTensorOp>(
          func, arg.getArgNumber(), syms);
  auto name = opToName[global];
  if (name.empty()) return nullptr;
  return SymbolRefAttr::get(func->getContext(), opToName[global]);
}

    for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
      SmallVector<TF::InitializeTableFromTextFileV2Op, 4>
          init_table_from_text_file_ops_to_erase;
      auto asset = LookupBoundInputOfType<AssetOp>(func, i, symbol_table);

      if (!asset) continue;

      auto arg = func.getArgument(i);
      bool arg_is_deletable = true;
      for (auto user : arg.getUsers()) {
        if (auto read_op =

    if (func::FuncOp func =
            dyn_cast<func::FuncOp>(barg.getOwner()->getParentOp())) {
      SymbolTable symbol_table(func->getParentOfType<ModuleOp>());
      auto global_tensor = tf_saved_model::LookupBoundInputOfType<
          tf_saved_model::GlobalTensorOp>(func, barg.getArgNumber(),
                                          symbol_table);
      ResourceConstructingOps result(global_tensor);
      return result;
    }

// refers to as a bound input, or null.
Operation *LookupBoundInput(func::FuncOp func, int arg_index,
                            const SymbolTable &symbol_table);

template <typename T>
T LookupBoundInputOfType(func::FuncOp func, int arg_index,
                         const SymbolTable &symbol_table) {
  return llvm::dyn_cast_or_null<T>(
      LookupBoundInput(func, arg_index, symbol_table));
}

    llvm::BitVector arg_indexes_to_remove(func.getNumArguments());

    // Replace arguments with const ops.
    for (BlockArgument argument : func.getArguments()) {
      auto asset = LookupBoundInputOfType<AssetOp>(
          func, argument.getArgNumber(), symbol_table);
      if (asset == nullptr) {
        continue;
      }

      // Create a const op for the asset if it doesn't already exist.

  SymbolTable symbol_table(module);
  for (auto func : module.getOps<func::FuncOp>()) {
    for (auto arg : func.getArguments()) {
      unsigned arg_number = arg.getArgNumber();
      auto global_tensor = LookupBoundInputOfType<GlobalTensorOp>(
          func, arg_number, symbol_table);
      if (!global_tensor) continue;

      auto arg_type = mlir::cast<RankedTensorType>(arg.getType());
      assert(arg_type.getRank() == 0);

    llvm::SmallVector<mlir::Type, 4> new_input_types;
    for (int i = 0, e = func.getNumArguments(); i < e; i++) {
      auto arg = func.getArgument(i);
      auto global_tensor = mlir::tf_saved_model::LookupBoundInputOfType<
          mlir::tf_saved_model::GlobalTensorOp>(func, i, symbol_table);
      if (global_tensor) {
        auto old_type = arg.getType();
        auto new_type =