#pragma once

#include "caffe2/core/common.h"
#include "caffe2/core/transform.h"
#include "caffe2/proto/caffe2_pb.h"
#include "caffe2/utils/proto_utils.h"

namespace caffe2 {

/**
 * Common Subexpression Elimination
 *
 * This transforms looks for specific operators (denoted by allowed_ops_),
 * and removes unnecessary repetition of that operator.
 *
 * Consider some operator of X, that reads from blob b_ written to by W.
 * X_a and X_b read the output of X. However, another operator Y, is the same
 * type as X, has the same arguments as X, and reads from the same input b_,
 * written to by W. It's output is the same as X. Y_a, Y_b, and Y_c read from Y.
 *
 * Then, we can eliminate the common subexpressions X and Y, and merge them to
 * Z, where X_a, X_b, Y_a, Y_b, and Y_c all read from Z.
 *
 *
 * TODO(benz): Fix the error to not match nodes that write to external output.
 */
class TORCH_API CommonSubexpressionEliminationTransform : public Transform {
 public:
  CommonSubexpressionEliminationTransform() {
    SetPatternMatchType(SORTED_WRT_EXECUTION_ORDER);
  }

 protected:
  bool PatternRule(
      const transform::Graph& g,
      const std::vector<int>& subgraph,
      int idx) override;
  bool ValidatorRule(
      const transform::Graph& g,
      const std::vector<int>& subgraph) override;
  bool ReplaceRule(const std::vector<int>& subgraph, transform::Graph* g_ptr)
      override;

 private:
  bool IsAllowed(string op_type) {
    return allowed_ops_.count(op_type);
  }
  std::set<string> allowed_ops_ = {"LearningRate", "FC"};
};

} // namespace caffe2