public:
  GraphPruningPass() = default;
  explicit GraphPruningPass(llvm::ArrayRef<std::string> ops_to_preserve);
  void runOnOperation() override;

 private:
  bool ShouldPreserveOp(Operation* op);
  bool ShouldPreserveIsland(IslandOp island);
  void PruneGraph(GraphOp graph);

  llvm::SmallDenseSet<mlir::StringAttr, 4> ops_to_preserve_ids_;
};

// Checks if a tf_executor.Graph can be pruned.