return func_arguments_;
  }
  const llvm::SmallVector<Value>& FuncOutputs() const { return func_outputs_; }

 private:
  // Compute once at construction and save as field.
  const llvm::SmallVector<Value> func_arguments_;
  const llvm::SmallVector<Value> func_outputs_;
};

// Helper data structure for output parameters to `ExtractSubgraphToFunc`.

  IRMapping values_in_scope;
  // Function arguments can appear as operands, so they clone should
  // be aware of them.
  assert(subgraph.FuncArguments().size() == new_func.getNumArguments());
  for (int i = 0; i < subgraph.FuncArguments().size(); ++i) {
    Value original_value = subgraph.FuncArguments()[i];
    Value new_func_arg = new_func.getArgument(i);
    values_in_scope.map(original_value, new_func_arg);
  }

            func::FuncOp func = while_op.body_function();
            auto func_argument = func.getArgument(next_op_operand_num);
            // Check if the op is returned without mutation. Returning values
            // from a while op follow return or identity -> return pattern.
            if (checkIfAnyUserIsConnectedToTermiantor(func_argument))
              next_values_to_visit.push_back(