bool IsFuncOpEmpty(func::FuncOp func_op) {
  return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}

// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
  if (IsFuncOpEmpty(func_op)) return {};

  auto graph_op_range = func_op.front().without_terminator();

  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.
// For TensorFlow V1.0 compatibility: when importing a graph without providing

// the fetch values in the main GraphOp corresponding to the original fetch
// values from `src_func_op`. Returns an empty vector when `src_func_op` is
// empty. `main_func_op` must have a GraphOp.
SmallVector<Value> CopyOpsToMainFunction(func::FuncOp src_func_op,
                                         func::FuncOp main_func_op) {
  GraphOp src_graph_op = GetGraphOpFromFuncOp(src_func_op);
  if (!src_graph_op) {

    // The pass runs on a valid tf_executor dialect, so the op should be the
    // GraphOp.
    return cast<GraphOp>(graph_op_range.begin());
  }

  return {};
}

void MergeDuplicateResourceOpsPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  GraphOp graph_op = GetGraphOpFromFuncOp(func_op);
  if (!graph_op) return;

  llvm::StringMap<Operation*> shared_name_to_resource;

// This pattern matches GraphOps with only one island, pulls out all inner ops
// of the island to the block containing the GraphOp, and then removes the
// GraphOp.
struct HoistInnerOpsSingleIslandGraph : public OpRewritePattern<GraphOp> {
  using OpRewritePattern<GraphOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(GraphOp op,
                                PatternRewriter &rewriter) const override {
    Block &block = op.GetBody();

    func.setAllResultAttrs(res_attrs);

    rewriter.setInsertionPointToStart(func.addEntryBlock());
    // In TFE, the function body is inlined in a GraphOp. Create a GraphOp
    // instance and move the regions from GraphFuncOp to GraphOp.
    auto executor_graph = rewriter.create<tf_executor::GraphOp>(
        loc, func.getFunctionType().getResults());

    // Replace the uses of block arguments with function arguments. Note that we

      std::next(copy_range.begin()) == copy_range.end() &&
      isa<tf_executor::GraphOp>(*copy_range.begin())) {
    // Already a graph.
    return;
  }

  auto return_op = dyn_cast<func::ReturnOp>(body.getTerminator());
  if (!return_op) {
    LLVM_DEBUG(llvm::dbgs() << "Expect function to end with return\n");
    return;
  }
  // Build GraphOp.
  OpBuilder builder(&body, body.begin());

    num_control_inputs_removed += island.getControlInputs().size();
    island.getControlInputsMutable().clear();
  } else if (auto fetch = dyn_cast<FetchOp>(op)) {
    GraphOp graph = fetch->getParentOfType<GraphOp>();
    int num_control_fetches = fetch.getNumOperands() - graph.getNumResults();
    if (num_control_fetches > 0) {
      fetch.getFetchesMutable().erase(graph.getNumResults(),

  MergedIsland* GetOperandCandidateToMergeWith(GraphOp graph,
                                               MergedIsland& merged_island);

  // Attempt to find an island group that consumes a result, either control or
  // data, from one of the islands in the given `merged_island`. If no candidate
  // can be found, returns nullptr.
  MergedIsland* GetResultCandidateToMergeWith(GraphOp graph,

  void runOnOperation() override;
};

// Extracts inner ops of tf_executor.island ops in a tf_executor.graph, in the
// order of ops in tf_executor.graph.
LogicalResult LiftIslandOpInnerOpsFromGraph(tf_executor::GraphOp graph) {
  auto graph_position = graph.getOperation()->getIterator();
  Block* parent_block = graph.getOperation()->getBlock();
  for (Operation& op : graph.GetBody().without_terminator()) {