raw_ostream& os() override { return string_stream; }

 private:
  std::string internal_str;
  llvm::raw_string_ostream string_stream;
};

}  // namespace

Status CreateFileForDumping(llvm::StringRef name,
                            std::unique_ptr<raw_ostream>* os,
                            std::string* filepath, llvm::StringRef dirname) {
  std::string dir;
  if (!dirname.empty())

void DataDumperLoggerConfig::DumpMlir(
    const std::string &filename,
    BridgeLoggerConfig::PrintCallbackFn print_callback) {
  std::unique_ptr<llvm::raw_ostream> os;
  std::string filepath;
  if (tensorflow::CreateFileForDumping(filename, &os, &filepath).ok()) {
    print_callback(*os);
    LOG(INFO) << "Dumped MLIR module to " << filepath;
  }
}

//
// This will create a file name via prefixing `name` with the value of the
// TF_DUMP_GRAPH_PREFIX environment variable if `dirname` is empty and
// suffixing `name` with ".mlir".
Status CreateFileForDumping(llvm::StringRef name,
                            std::unique_ptr<llvm::raw_ostream>* os,
                            std::string* filepath,
                            llvm::StringRef dirname = "");

  std::string name = llvm::formatv("{0,0+4}_mlir_bridge_{1}_{2}", log_counter++,
                                   pass_name, file_suffix);

  std::unique_ptr<llvm::raw_ostream> os;
  std::string filepath;
  if (CreateFileForDumping(name, &os, &filepath).ok()) {
    print_callback(*os);
    LOG(INFO) << "Dumped MLIR module to " << filepath;
  }
}

void BridgeLoggerConfig::printBeforeIfEnabled(mlir::Pass* pass,