#pragma once

#include <torch/csrc/WindowsTorchApiMacro.h>

#include <torch/csrc/jit/codegen/cuda/dispatch.h>
#include <torch/csrc/jit/codegen/cuda/ir_base_nodes.h>

#include <unordered_map>

namespace torch {
namespace jit {
namespace fuser {

struct Fusion;

/*
 * Mutators are the mechanism used to modify IR nodes. Since most nodes are
 * immutable or at least partially immutable changeing them can require creating
 * a new node. Base mutator at the moment is a dumb sample mutator that takes
 * any float of value 1.0 and converts it to 0.0; It is currently used as a
 * dummy example, however, we should make it a simple instantiation of all the
 * mutate functions on all node types so that people can inhereit it, and only
 * specialize those nodes which they want to have a particular transformation.
 */

// Search through "within" and replace all instances of "instance" with the
// value "with".
struct TORCH_CUDA_API ReplaceAll : public OptOutMutator {
 private:
  // Will look in fusion and if we're replacing an input or output we register
  // those changes
  void replaceInpOut();

  ReplaceAll(Val* _instance, Val* _with) {
    registerMutation(_instance, _with);
  }
  ReplaceAll(std::unordered_map<Val*, Val*> _replacement_map) {
    for (auto it : _replacement_map)
      registerMutation(it.first, it.second);
  }

 public:
  static void instancesOf(Val* instance, Val* with);
  static void instancesOf(std::unordered_map<Val*, Val*> replacement_map);
  static void instancesWithin(Val* instance, Val* with, Expr* within);
  static void instancesWithin(
      std::unordered_map<Val*, Val*> replacement_map,
      Expr* within);
};

} // namespace fuser
} // namespace jit
} // namespace torch
