}

};  // namespace

LogicalResult Tf2XlaRewriter::RewriteOp(Operation* op,
                                        PatternRewriter& rewriter,
                                        const std::string& device_type) {
  Tf2XlaRewriter tf2xla_rewriter(op, rewriter, device_type);
  return tf2xla_rewriter.LegalizeOp();
}

Tf2XlaRewriter::Tf2XlaRewriter(Operation* op, PatternRewriter& rewriter,

namespace mlir {
namespace mhlo {

class Tf2XlaRewriterTestPeer;

class Tf2XlaRewriter {
 public:
  static mlir::LogicalResult RewriteOp(mlir::Operation* op,
                                       mlir::PatternRewriter& rewriter,
                                       const std::string& device_type);

 private:
  friend class Tf2XlaRewriterTestPeer;

  Tf2XlaRewriter(mlir::Operation* op, mlir::PatternRewriter& rewriter,

      XlaComputation& computation) {
    return tf2xla_rewriter_.ImportXlaComputation(computation);
  }

 private:
  OpBuilder op_builder_;
  EmptyPatternRewriter empty_rewriter_;
  Tf2XlaRewriter tf2xla_rewriter_;
};

// This should only have unit tests. End to end tests should be done with
// FILECHECK and MLIR tests.
class Tf2XlaRewriterTest : public ::testing::Test {
 public:

    if (!(IsOpAllowedTf2xlaFallback(abstractOp->getTypeID()) ||
          (prefer_tf2xla_ &&
           IsOpAllowedTf2xlaPreferred(abstractOp->getTypeID())))) {
      return failure();
    }

    return Tf2XlaRewriter::RewriteOp(op, rewriter, device_type_);
  }

 private:
  std::string device_type_;
  bool prefer_tf2xla_;
  bool use_tf2xla_hlo_importer_;
};