// Part of the Concrete Compiler Project, under the BSD3 License with Zama
// Exceptions. See
// https://github.com/zama-ai/concrete/blob/main/LICENSE.txt
// for license information.

// Infrastructure for type inference
//
// The type inference framework below allows for the implementation of
// custom type inference passes that propagate the types of a
// partially-typed input IR throughout its operations following the
// flow of data.
//
// The framework relies on three main concepts:
//
// - A type resolver inheriting from `TypeResolver` that (1) specifies
//   which types are to be considered as unresolved and that (2)
//   resolves the actual types for the values related to an operation
//   based on a previous state of type inference.
//
// - The propagation of types is implemented through the two classes
//   `ForwardTypeInferenceAnalysis` and
//   `BackwardTypeInferenceAnalysis`, which can be used as forward and
//   backward dataflow analyses with the MLIR sparse dataflow analysis
//   framework. These analyses invoke the above-mentioned type
//   resolver for an operation every time the inferred type of a value
//   referenced by that operation changes.
//
// - The type inference state for an operation is represented by an
//   instance of the class `LocalInferenceState`, which maps the
//   values related to an operation to instances of `InferredType`
//   (either indicating the inferred type as an `mlir::Type` or
//   indicating that no type has been inferred, yet).
//
// Additionally, the local rules specifying the relationship between
// the state of inference before invocation of the type resolver and
// the state of inference the invocation can be implemented with type
// constraints (instances of sub-classes of `TypeConstraint`). Type
// constraints can be combined into a `TypeConstraintSet`, which
// provides a method `converge()`, that attempts to apply the
// constraints until the resulting type inference state converges.
//
// There are multiple, predefined type constraint classes for common
// constraints (e.g., if two values must have the same type or the
// same element type). These exist both as static constraints and as
// dynamic constraints. Some pre-defined type constraints depend on a
// class that yields a pair of values for which the constraints shall
// be applied (e.g., yielding two operands or an operand and a result,
// etc.).
//
// The global state of type inference after running the type inference
// analyses is contained in the `DataFlowSolver` to which the analyses
// where added. The local state of inference for an operation can be
// obtained at any stage of the analysis via the helper methods of
// `TypeInferenceUtils`.

#ifndef CONCRETELANG_ANALYSIS_TYPEINFERENCEANALYSIS_H
#define CONCRETELANG_ANALYSIS_TYPEINFERENCEANALYSIS_H

#include <concretelang/Dialect/TypeInference/IR/TypeInferenceOps.h>
#include <llvm/ADT/TypeSwitch.h>
#include <llvm/Support/Debug.h>
#include <mlir/Analysis/DataFlow/SparseAnalysis.h>
#include <mlir/Dialect/Func/IR/FuncOps.h>

namespace mlir {
namespace concretelang {
// Encapsulated the state of inference for a single type
class InferredType {
public:
  InferredType() : type(std::nullopt) {}
  InferredType(std::optional<mlir::Type> t) : type(t) {}
  InferredType(mlir::Type t) : type(t) {}

  mlir::Type getType() const { return type.value(); }
  void setType(std::optional<mlir::Type> t) { type = t; }
  bool hasType() const { return type.has_value(); }

  bool operator==(const InferredType &other) const {
    return other.type == type;
  }
  bool operator!=(const InferredType &other) const {
    return other.type != type;
  }

protected:
  std::optional<mlir::Type> type;
};

inline raw_ostream &operator<<(llvm::raw_ostream &os, InferredType t) {
  if (t.hasType())
    os << t.getType();
  else
    os << "(none)";

  return os;
}

// A local inference state is a mapping of values to inferred types
// used to represent the state of type inference for all values
// related to an operation (e.g., types inferred for operands,
// results, block arguments, etc.)
class LocalInferenceState {
public:
  using MapT = llvm::DenseMap<mlir::Value, InferredType>;

  LocalInferenceState() : changed(false) {}

  // Updates the inferred type for a value
  void set(mlir::Value v, InferredType t) {
    if (find(v) != t)
      changed = true;

    state[v] = t;
  }

  // Looks up the inferred type as an `InferredType` for a value. If
  // no type has been inferred or set for the value, an empty inferred
  // type is returned.
  InferredType find(mlir::Value v) const {
    auto it = state.find(v);

    return (it != state.end()) ? it->second : InferredType(std::nullopt);
  }

  // Looks up the inferred type as an MLIR type for a value. If no
  // type has been inferred or set for the value, an empty type is
  // returned.
  mlir::Type lookup(mlir::Value v) const {
    InferredType t = this->find(v);

    return t.hasType() ? t.getType() : mlir::Type();
  }

  // Looks up the inferred types for an entire range of values. For
  // values whose type has not been inferred, empty types are
  // returned.
  llvm::SmallVector<mlir::Type> lookup(mlir::ValueRange r) const {
    return llvm::to_vector(
        llvm::map_range(r, [&](mlir::Value v) { return this->lookup(v); }));
  }

  // Resets the change flag
  void setUnchanged() { changed = false; }

  // Indicates if the inferred type for any of the values has changed
  // or if the inferred type for a new value has been added since the
  // creation of the inference set or since the last time the change
  // flag was reset.
  bool hasChanged() { return changed; }

  // Helper class representing the values for which types have been
  // inferred
  class ValueSet {
  public:
    ValueSet(const MapT &map) : map(map) {}

    template <typename MapIteratorT> class IteratorImpl {
    public:
      IteratorImpl(MapIteratorT it) : it(it) {}
      mlir::Value operator*() { return it->first; }
      mlir::Value *operator->() { return &it->first; }
      IteratorImpl operator++() { return IteratorImpl(it++); }
      bool operator==(const IteratorImpl &other) { return other.it == it; }
      bool operator!=(const IteratorImpl &other) { return other.it != it; }

    protected:
      MapIteratorT it;
    };

    using const_iterator = IteratorImpl<MapT::const_iterator>;

    const_iterator cbegin() const { return const_iterator(map.begin()); }
    const_iterator cend() const { return const_iterator(map.end()); }
    const_iterator begin() const { return const_iterator(map.begin()); }
    const_iterator end() const { return const_iterator(map.end()); }

  protected:
    const MapT &map;
  };

  void dump() const {
    llvm::dbgs() << "LocalInferenceState:\n";

    for (auto [v, t] : state) {
      llvm::dbgs() << "  " << v << ": " << t << "\n";
    }
  }

  const ValueSet getValues() const { return ValueSet(state); }

protected:
  MapT state;
  bool changed;
};

// A type resolver is responsible for the actual inference of the
// types related to an operation based on the previous inference state
class TypeResolver {
public:
  virtual ~TypeResolver() {}

  // Resolve the types for all values related to `op` based on the
  // previous inference state `prevState` and the current state
  // `currState`.
  virtual LocalInferenceState resolve(mlir::Operation *op,
                                      const LocalInferenceState &prevState) = 0;

  // Returns true for any type that is considered unresolved by the
  // resolver
  virtual bool isUnresolvedType(mlir::Type t) const = 0;

  // Invokes the callback function `fn` to all values related to the
  // operation `op`. Iteration over the related value stops if `fn`
  // returns `false`.
  //
  // Returns `true` if `fn` returned `true` for all related values,
  // otherwise `false`.
  virtual bool
  iterateRelatedValues(mlir::Operation *op,
                       llvm::function_ref<bool(mlir::Value)> fn) const {
    auto applyToOperandsAndResults = [&](mlir::Operation *op) {
      for (mlir::Value v : op->getOperands()) {
        if (!fn(v))
          return false;
      }

      for (mlir::Value v : op->getResults()) {
        if (!fn(v))
          return false;
      }

      return true;
    };

    if (!applyToOperandsAndResults(op))
      return false;

    for (mlir::Region &r : op->getRegions()) {
      for (mlir::Value v : r.getArguments()) {
        if (!fn(v))
          return false;
      }

      if (!op->hasTrait<mlir::OpTrait::NoTerminator>()) {
        // Also map return-like terminators, as they are excluded from
        // forward and backward analysis and not visited individually
        for (mlir::Block &b : r.getBlocks()) {
          mlir::Operation *terminator = b.getTerminator();

          if (terminator->hasTrait<OpTrait::ReturnLike>())
            if (!applyToOperandsAndResults(terminator))
              return false;
        }
      }
    }

    return true;
  }
};

// Lattice value for type inference analysis, encapsulating a single
// inferred type
class TypeInferenceLatticeValue {
public:
  TypeInferenceLatticeValue(const InferredType &t = {}) : inferredType(t) {}

  static TypeInferenceLatticeValue getPessimisticValueState(mlir::Value value) {
    return TypeInferenceLatticeValue();
  }

  bool operator==(const TypeInferenceLatticeValue &rhs) {
    return rhs.getInferredType() == inferredType;
  }

  static TypeInferenceLatticeValue join(const TypeInferenceLatticeValue &lhs,
                                        const TypeInferenceLatticeValue &rhs) {
    return TypeInferenceLatticeValue(lhs.getInferredType().hasType()
                                         ? lhs.getInferredType()
                                         : rhs.getInferredType());
  }

  static TypeInferenceLatticeValue meet(const TypeInferenceLatticeValue &lhs,
                                        const TypeInferenceLatticeValue &rhs) {
    return join(lhs, rhs);
  }

  const InferredType &getInferredType() const { return inferredType; }
  void setType(const InferredType &t) { this->inferredType = t; }

  void print(raw_ostream &os) const { os << inferredType; }

protected:
  InferredType inferredType;
};

// Lattice for forward and backward type inference analysis
class TypeInferenceLattice
    : public mlir::dataflow::Lattice<TypeInferenceLatticeValue> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TypeInferenceLattice)

  using Lattice::Lattice;
};

// Collection of common operations for forward and backward type
// inference and subsequent transformations (e.g., rewriting the IR
// with the inferred types).
class TypeInferenceUtils {
public:
  // Looks up the inferred types for all values related to the
  // operation `op`, i.e., the inferred types for operands, results,
  // region arguments and operands of terminators used in directly
  // nested regions.
  static LocalInferenceState getLocalInferenceState(
      mlir::Operation *op, const TypeResolver &typeResolver,
      llvm::function_ref<const TypeInferenceLattice *(mlir::Value)> lookup) {
    LocalInferenceState map;

    typeResolver.iterateRelatedValues(op, [&](mlir::Value v) {
      map.set(v, getInferredType(v, lookup));
      return true;
    });

    return map;
  }

  // Retrieves the inferred types for all values related to `op` using
  // `DataFlowSolver::lookupState` invoked on `solver`.
  static LocalInferenceState
  getLocalInferenceState(const DataFlowSolver &solver,
                         const TypeResolver &typeResolver,
                         mlir::Operation *op) {
    LocalInferenceState map = TypeInferenceUtils::getLocalInferenceState(
        op, typeResolver, [&](mlir::Value v) {
          return solver.lookupState<TypeInferenceLattice>(v);
        });

    return map;
  }

  // If `t` is a tensor or memref type, return a tensor or memref type
  // with the same shape, but with `elementType` as the element
  // type. If `t` is a scalar type, simply return `elementType`.
  static mlir::Type applyElementType(mlir::Type elementType, mlir::Type t) {
    if (mlir::RankedTensorType rtt =
            llvm::dyn_cast<mlir::RankedTensorType>(t)) {
      return mlir::RankedTensorType::get(
          rtt.getShape(), applyElementType(elementType, rtt.getElementType()));
    } else if (mlir::MemRefType mrt = llvm::dyn_cast<mlir::MemRefType>(t)) {
      return mlir::MemRefType::get(
          mrt.getShape(), applyElementType(elementType, mrt.getElementType()));
    } else {
      return elementType;
    }
  }

protected:
  // Looks up the lattice value for `v` using `lookup` and returns the
  // result as an instance of `InferredType`
  static InferredType getInferredType(
      mlir::Value v,
      llvm::function_ref<const TypeInferenceLattice *(mlir::Value)> lookup) {
    const TypeInferenceLattice *latticeOperand = lookup(v);

    if (latticeOperand) {
      return InferredType(latticeOperand->getValue().getInferredType());
    } else {
      return InferredType{};
    }
  }
};

// Base class for all type constraints that can be added to a
// `TypeConstraintSet`.
class TypeConstraint {
public:
  virtual ~TypeConstraint() {}

  // Apply the constraint on the types inferred previously (based on
  // types inferred during previous visits of related operations or
  // the operation itself) and those inferred during the convergence
  // of the current visit of the operation.
  virtual void apply(mlir::Operation *op, TypeResolver &resolver,
                     LocalInferenceState &nextState,
                     const LocalInferenceState &prevState) = 0;

protected:
  // Returns the type inferred for `v`. This is either the type of `v`
  // if `v` has a non-unresolved type, or the type specified in
  // `currState`, or the type specified in `prevState`, or
  // `std::nullopt` if the type of `v` is unresolved and no type has
  // been inferred.
  static InferredType getFirstTypeInOrder(TypeResolver &resolver,
                                          const LocalInferenceState &currState,
                                          const LocalInferenceState &prevState,
                                          mlir::Value v) {
    if (!resolver.isUnresolvedType(v.getType()))
      return v.getType();

    InferredType sit = currState.find(v);

    if (sit.hasType())
      return sit;

    return prevState.find(v);
  }
};

// Yields the operands with the static indices i and j of an operation
template <int i, int j> class OperandValueYield {
public:
  static std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {op->getOperand(i), op->getOperand(j)};
  }
};

// Yields the results with the static indices i and j of an
// operation
template <int i, int j> class ResultValueYield {
public:
  static std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {op->getResult(i), op->getResult(j)};
  }
};

// Yields the operand and result with the static indices i and j of an
// operation
template <int i, int j> class OperandAndResultValueYield {
public:
  static std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {op->getOperand(i), op->getResult(j)};
  }
};

// Yields the i-th and j-th operands of an operation, where `i` and
// `j` are defined dynamically
class DynamicOperandValueYield {
public:
  DynamicOperandValueYield(int i, int j) : i(i), j(j) {}

  std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {op->getOperand(i), op->getOperand(j)};
  }

protected:
  int i;
  int j;
};

// Yields the i-th and j-th results of an operation, where `i` and
// `j` are defined dynamically
class DynamicResultValueYield {
public:
  DynamicResultValueYield(int i, int j) : i(i), j(j) {}

  std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {op->getResult(i), op->getResult(j)};
  }

protected:
  int i;
  int j;
};

// Yields the i-th operand and j-th result of an operation, where `i`
// and `j` are defined dynamically
class DynamicOperandAndResultValueYield {
public:
  DynamicOperandAndResultValueYield(int i, int j) : i(i), j(j) {}

  std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {op->getOperand(i), op->getResult(j)};
  }

protected:
  int i;
  int j;
};

// Yield the return values of two callback functions `f̀` and `g`
class DynamicFunctorYield {
public:
  DynamicFunctorYield(std::function<mlir::Value(void)> f,
                      std::function<mlir::Value(void)> g)
      : f(f), g(g) {}

  std::pair<mlir::Value, mlir::Value> yield(mlir::Operation *op) {
    return {f(), g()};
  }

protected:
  std::function<mlir::Value()> f;
  std::function<mlir::Value()> g;
};

// Empty type constraint that accepts any previously inferred type
class NoTypeConstraint : public TypeConstraint {
public:
  void apply(mlir::Operation *op, TypeResolver &resolver,
             LocalInferenceState &currState,
             const LocalInferenceState &prevState) override {
    for (mlir::Value v : prevState.getValues()) {
      InferredType t = getFirstTypeInOrder(resolver, currState, prevState, v);
      currState.set(v, t);
    }
  }
};

// Base class for statically-defined type constraints, assuming `YieldT::yield`
// is a static function
template <typename YieldT> class StaticTypeConstraint : public TypeConstraint {
public:
  StaticTypeConstraint() = default;

protected:
  std::pair<mlir::Value, mlir::Value> yieldValues(mlir::Operation *op) {
    return YieldT::yield(op);
  }
};

// Base class for dynamically-defined type constraints, assuming `YieldT::yield`
// is a non-static function
template <typename YieldT> class DynamicTypeConstraint : public TypeConstraint {
public:
  template <typename... YieldTCtorArgTs>
  DynamicTypeConstraint(YieldTCtorArgTs &&...ctorArgs)
      : yield(std::forward<YieldTCtorArgTs>(ctorArgs)...) {}

protected:
  std::pair<mlir::Value, mlir::Value> yieldValues(mlir::Operation *op) {
    return yield.yield(op);
  }

  YieldT yield;
};

// Constraint that ensures that the types inferred for two values
// match exactly. If different types have been inferred previously for
// the two values, precedence is given to the type inferred for the
// first value yielded by `BaseConstraintT::yield`.
//
// If the types for the two values are fixed, but contradictory, an
// assertion is triggered.
template <typename BaseConstraintT>
class SameTypeConstraintBase : public BaseConstraintT {
public:
  using BaseConstraintT::BaseConstraintT;

  void apply(mlir::Operation *op, TypeResolver &resolver,
             LocalInferenceState &currState,
             const LocalInferenceState &prevState) override {
    auto [leftValue, rightValue] = this->yieldValues(op);

    mlir::Type leftType = leftValue.getType();
    mlir::Type rightType = rightValue.getType();

    bool leftUnresolved = resolver.isUnresolvedType(leftType);
    bool rightUnresolved = resolver.isUnresolvedType(rightType);

    // Priority is given from left to right, i.e., the type of the
    // first value yielded by `yield()` takes precedence over the
    // type for the second value
    if (!leftUnresolved && rightUnresolved) {
      currState.set(leftValue, leftType);
      currState.set(rightValue, leftType);
    } else if (leftUnresolved && !rightUnresolved) {
      currState.set(leftValue, rightType);
      currState.set(rightValue, rightType);
    } else if (!leftUnresolved && !rightUnresolved) {
      assert(leftType == rightType && "Constraint cannot be matched, as "
                                      "operands have different, fixed types");
      currState.set(leftValue, leftType);
      currState.set(rightValue, rightType);
    } else {
      // Both unresolved
      InferredType leftCurrType =
          this->getFirstTypeInOrder(resolver, currState, prevState, leftValue);
      InferredType rightCurrType =
          this->getFirstTypeInOrder(resolver, currState, prevState, rightValue);

      if (leftCurrType.hasType()) {
        currState.set(leftValue, leftCurrType);
        currState.set(rightValue, leftCurrType);
      } else if (rightCurrType.hasType()) {
        currState.set(leftValue, rightCurrType);
        currState.set(rightValue, rightCurrType);
      }
    }
  }
};

// Same type constraint for static yield operators
template <typename YieldT>
using SameTypeConstraint = SameTypeConstraintBase<StaticTypeConstraint<YieldT>>;

// Same type constraint for dynamic yield operators
template <typename YieldT>
using DynamicSameTypeConstraint =
    SameTypeConstraintBase<DynamicTypeConstraint<YieldT>>;

// Base for constraints that ensure that the types inferred for two
// values have the same nested type. What is considered a nested type
// (e.g., an arbitrarily deeply nested scalar type or a shallow type
// of a pointer) can be customized through `getNestedType()` and
// `applyNestedType()`.
//
// The depth of the nesting for the left and right values can be set
// independently through the template parameters `leftDepth` and
// `rightDepth`.
//
// If different types have been inferred previously for the two
// values, precedence is given to the element type inferred for the
// first ("left") value yielded by `BaseConstraintT::yield`.
//
// If the element types for the two values are fixed, but
// contradictory, an assertion is triggered.
template <typename BaseConstraintT, unsigned int leftDepth = 1,
          unsigned int rightDepth = 1>
class SameNestedTypeConstraintBase : public BaseConstraintT {
public:
  using BaseConstraintT::BaseConstraintT;

  void apply(mlir::Operation *op, TypeResolver &resolver,
             LocalInferenceState &currState,
             const LocalInferenceState &prevState) override {
    auto [leftValue, rightValue] = this->yieldValues(op);

    mlir::Type leftType = leftValue.getType();
    mlir::Type rightType = rightValue.getType();

    bool leftUnresolved = resolver.isUnresolvedType(leftType);
    bool rightUnresolved = resolver.isUnresolvedType(rightType);

    // Priority is given from left to right, i.e., the type of the
    // first value yielded by `yield()` takes precedence over the
    // type for the second value
    if (!leftUnresolved && rightUnresolved) {
      currState.set(leftValue, leftType);
      currState.set(rightValue,
                    applyNestedType(getNestedType(leftType, leftDepth),
                                    rightType, rightDepth));
    } else if (leftUnresolved && !rightUnresolved) {
      currState.set(leftValue,
                    applyNestedType(getNestedType(rightType, rightDepth),
                                    leftType, leftDepth));
      currState.set(rightValue, rightType);
    } else if (!leftUnresolved && !rightUnresolved) {
      assert(getNestedType(leftType, leftDepth) ==
                 getNestedType(rightType, rightDepth) &&
             "Constraint cannot be matched, as fixed types of the values"
             "are incompatible types");
      currState.set(leftValue, leftType);
      currState.set(rightValue, rightType);
    } else {
      // Both unresolved
      InferredType leftCurrType =
          this->getFirstTypeInOrder(resolver, currState, prevState, leftValue);
      InferredType rightCurrType =
          this->getFirstTypeInOrder(resolver, currState, prevState, rightValue);

      if (leftCurrType.hasType()) {
        currState.set(leftValue, leftCurrType);
        currState.set(
            rightValue,
            applyNestedType(getNestedType(leftCurrType.getType(), leftDepth),
                            rightType, rightDepth));
      } else if (rightCurrType.hasType()) {
        currState.set(
            leftValue,
            applyNestedType(getNestedType(rightCurrType.getType(), rightDepth),
                            leftType, leftDepth));
        currState.set(rightValue, rightCurrType);
      }
    }
  }

protected:
  virtual mlir::Type getNestedType(mlir::Type t, unsigned int depth) {
    mlir::Type res = t;

    for (unsigned int i = 0; i < depth; i++)
      res = getNestedType(res);

    return res;
  }

  virtual mlir::Type getNestedType(mlir::Type t) = 0;

  virtual mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t,
                                     unsigned int depth) {
    if (depth == 0) {
      return nestedType;
    } else {
      return applyNestedType(
          applyNestedType(nestedType, getNestedType(t), depth - 1), t);
    }
  }

  virtual mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) = 0;
};

// Constraint that ensures that the types inferred for two values use
// the same element type. If different types have been inferred
// previously for the two values, precedence is given to the element
// type inferred for the first value yielded by
// `BaseConstraintT::yield`.
//
// If the element types for the two values are fixed, but
// contradictory, an assertion is triggered.
template <typename BaseConstraintT>
class SameElementTypeConstraintBase
    : public SameNestedTypeConstraintBase<BaseConstraintT> {
public:
  using SameNestedTypeConstraintBase<
      BaseConstraintT>::SameNestedTypeConstraintBase;

protected:
  mlir::Type getNestedType(mlir::Type t) override {
    if (mlir::ShapedType sht = llvm::dyn_cast<mlir::ShapedType>(t)) {
      return sht.getElementType();
    } else {
      return t;
    }
  }

  mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) override {
    return TypeInferenceUtils::applyElementType(nestedType, t);
  }
};

// Same element type constraint for static yield operators
template <typename YieldT>
using SameElementTypeConstraint =
    SameElementTypeConstraintBase<StaticTypeConstraint<YieldT>>;

// Same element type constraint for dynamic yield operators
template <typename YieldT>
using DynamicSameElementTypeConstraint =
    SameElementTypeConstraintBase<DynamicTypeConstraint<YieldT>>;

// Type constraint ensuring that two operands with the statically-defined
// indexes `a` and `b` of an operation have the same type
template <int a, int b>
using SameOperandTypeConstraint = SameTypeConstraint<OperandValueYield<a, b>>;

// Type constraint ensuring that two results with the statically-defined
// indexes `a` and `b` of an operation have the same type
template <int a, int b>
using SameResultTypeConstraint = SameTypeConstraint<ResultValueYield<a, b>>;

// Type constraint ensuring that the operand with the
// statically-defined index `a` and the result with the
// statically-defined index `b` of an operation have the same type
template <int operandIdx, int resultIdx>
using SameOperandAndResultTypeConstraint =
    SameTypeConstraint<OperandAndResultValueYield<operandIdx, resultIdx>>;

// Type constraint ensuring that two operands with the statically-defined
// indexes `a` and `b` of an operation have the same element type
template <int a, int b>
using SameOperandElementTypeConstraint =
    SameElementTypeConstraint<OperandValueYield<a, b>>;

// Type constraint ensuring that two results with the statically-defined
// indexes `a` and `b` of an operation have the same element type
template <int a, int b>
using SameResultElementTypeConstraint =
    SameElementTypeConstraint<ResultValueYield<a, b>>;

// Type constraint ensuring that the operand with the
// statically-defined index `a` and the result with the
// statically-defined index `b` of an operation have the same element type
template <int operandIdx, int resultIdx>
using SameOperandAndResultElementTypeConstraint = SameElementTypeConstraint<
    OperandAndResultValueYield<operandIdx, resultIdx>>;

namespace {
namespace impl {
// Specialization needs to happen at namespace scope
template <typename ConstraintT, typename... ArgTs>
void addConstraint(std::vector<std::unique_ptr<TypeConstraint>> &constraints,
                   ArgTs &&...args) {
  std::unique_ptr<TypeConstraint> constraint =
      std::make_unique<ConstraintT>(std::forward<ArgTs>(args)...);
  constraints.push_back(std::move(constraint));
}

template <typename ConstraintT, typename... ArgTs>
void addConstraints(std::vector<std::unique_ptr<TypeConstraint>> &constraints,
                    ArgTs &&...args) {
  addConstraint<ConstraintT>(constraints, std::forward<ArgTs>(args)...);
}

template <typename ConstraintT0, typename ConstraintT1, typename... TailTs,
          typename... ArgTs>
void addConstraints(std::vector<std::unique_ptr<TypeConstraint>> &constraints,
                    ArgTs &&...args) {
  addConstraint<ConstraintT0>(constraints, std::forward<ArgTs>(args)...);
  addConstraints<ConstraintT1, TailTs..., ArgTs...>(
      constraints, std::forward<ArgTs>(args)...);
}

}; // namespace impl
}; // namespace

// Set of type constraints to be applied on one visit of an
// operation. The template parameter `maxApplications` sets a limit to
// the number of rounds the set of constraints is applied in
// `converge()`.
template <int maxApplications = 10> class TypeConstraintSet {
public:
  TypeConstraintSet() {}

  // Instantiates a constraint of the type `ConstraintT` with the
  // arguments `args` and adds the constraint to the set
  template <typename ConstraintT, typename... ArgTs>
  void addConstraint(ArgTs &&...args) {
    impl::addConstraint<ConstraintT, ArgTs...>(constraints,
                                               std::forward<ArgTs>(args)...);
  }

  // Instantiates the constraints of types specified by `ConstraintTs` with the
  // arguments `args` and adds them to the set
  template <typename... ConstraintTs, typename... ArgTs>
  void addConstraints(ArgTs &&...args) {
    impl::addConstraints<ConstraintTs..., ArgTs...>(
        constraints, std::forward<ArgTs>(args)...);
  }

  // Applies all type constraints in order of their addition for a
  // maximum of `maxApplications` rounds. Initial state needs to be
  // provided in `prevState`. The resulting state is contained in
  // `currState` afterward th call.
  //
  // If the rules converge on less or equal to `maxApplications`
  // rounds, the function return `true`, otherwise `false`.
  bool converge(mlir::Operation *op, TypeResolver &resolver,
                LocalInferenceState &currState,
                const LocalInferenceState &prevState) {
    // Initialize with fixed types
    for (mlir::Value v : op->getOperands()) {
      if (!resolver.isUnresolvedType(v.getType()))
        currState.set(v, v.getType());
    }

    for (mlir::Value v : op->getResults()) {
      if (!resolver.isUnresolvedType(v.getType()))
        currState.set(v, v.getType());
    }

    // Apply type constraints until state converges or maximum number
    // of iterations has been reached
    for (int i = 0; i < maxApplications; i++) {
      currState.setUnchanged();

      for (auto &constraint : constraints) {
        constraint->apply(op, resolver, currState, prevState);
      }

      if (!currState.hasChanged())
        return true;
    }

    return !currState.hasChanged();
  }

protected:
  std::vector<std::unique_ptr<TypeConstraint>> constraints;
};

// Base class for forward and backward type analysis with shared
// infrastructure. The template parameter `AnalysisT` must either be
// `mlir::dataflow::SparseDataFlowAnalysis` or
// `mlir::dataflow::SparseBackwardDataFlowAnalysis`.
template <typename AnalysisT>
class TypeInferenceAnalysisBase : public AnalysisT {
public:
  // Constructor passing on `constructorArgs` to the analysis class
  template <typename... ConstructorArgTs>
  TypeInferenceAnalysisBase(TypeResolver &resolver,
                            mlir::DataFlowSolver &solver,
                            ConstructorArgTs &&...constructorArgs)
      : AnalysisT(solver, std::forward<ConstructorArgTs>(constructorArgs)...),
        resolver(resolver) {}

protected:
  // Returns the types inferred for all values related to the
  // operation `op` from the respective lattice values. For values,
  // for which no type has been inferred so far, a new lattice with an
  // empty inferred types is created.
  LocalInferenceState getCurrentInferredTypes(mlir::Operation *op) {
    LocalInferenceState map = TypeInferenceUtils::getLocalInferenceState(
        op, resolver, [&](mlir::Value v) {
          return this->template getOrCreate<TypeInferenceLattice>(v);
        });

    return map;
  }

  // Dumps the inferred type for all values related to `op` on
  // `llvm::dbgs()`.
  void debugPrintOp(mlir::Operation *op) {
    LocalInferenceState inferredTypes = getCurrentInferredTypes(op);

    llvm::dbgs() << "Inference state for " << op->getName() << " \n";
    for (mlir::Value operand : op->getOperands()) {
      InferredType operandType = inferredTypes.find(operand);
      llvm::dbgs() << "  Operand: " << operandType << "\n";
    }

    for (mlir::Value result : op->getResults()) {
      InferredType resultType = inferredTypes.find(result);
      llvm::dbgs() << "  Result: " << resultType << "\n";
    }

    for (mlir::Region &r : op->getRegions()) {
      for (mlir::Value v : r.getArguments()) {
        InferredType argType = inferredTypes.find(v);
        llvm::dbgs() << "  RegionArg: " << argType << "\n";
      }
    }
  }

  // Updates the lattice for the inferred type for the value `v` using
  // its inferred type specified by `state`.
  void updateLatticeValuesAndDependenciesFromState(LocalInferenceState &state,
                                                   mlir::Value v,
                                                   mlir::Operation *op) {
    TypeInferenceLattice *lattice =
        this->template getOrCreate<TypeInferenceLattice>(v);

    this->addDependency(lattice, op);

    auto latticeValue = lattice->getValue();
    latticeValue.setType(state.find(v));
    mlir::ChangeResult res = lattice->join(latticeValue);
    this->propagateIfChanged(lattice, res);
  }

  // Visits all operations recursively by invoking `doVisitOperation`
  // on them. Since `doVisitOperation` registers dependencies between
  // related values an operations explicitly, this ensures that all
  // operations are visited upon an update of any of their
  // dependencies, even those for which the dataflow analysis
  // framework does invoke `visitOperation`.
  void initializeRecursively(mlir::Operation *op) {
    for (Region &region : op->getRegions()) {
      for (Operation &op : region.getOps()) {
        doVisitOperation(&op);
        initializeRecursively(&op);
      }
    }
  }

  LogicalResult initialize(mlir::Operation *op) override {
    initializeRecursively(op);

    return AnalysisT::initialize(op);
  }

  // Visits a single operation: looks up the lattice values for the
  // values related to the operation, extracts the inferred types,
  // invokes the type resolver and updates the lattice values with the
  // newly inferred types.
  void doVisitOperation(Operation *op) {
    // Skip operations that do not use unresolved types
    if (resolver.iterateRelatedValues(op, [&](mlir::Value v) {
          return !resolver.isUnresolvedType(v.getType());
        })) {
      return;
    }

    // Retrieve the current state of inference from the lattice values
    const LocalInferenceState inferredTypes = getCurrentInferredTypes(op);
    LocalInferenceState state;

    // Treat special op types for type inference
    mlir::TypeSwitch<mlir::Operation *>(op)
        .Case<TypeInference::PropagateUpwardOp>([&](auto op) {
          state.set(op->getOperand(0), op->getResult(0).getType());
          state.set(op->getResult(0), op->getResult(0).getType());
        })
        .template Case<TypeInference::PropagateDownwardOp>([&](auto op) {
          state.set(op->getOperand(0), op->getOperand(0).getType());
          state.set(op->getResult(0), op->getOperand(0).getType());
        })

        // Actually resolve the types based on the current state of
        // inference
        .Default([&](auto op) { state = resolver.resolve(op, inferredTypes); });

    // Store the resulting state in the lattice values
    resolver.iterateRelatedValues(op, [&](mlir::Value v) {
      updateLatticeValuesAndDependenciesFromState(state, v, op);
      return true;
    });
  }

  // Redirects the request for the initial inferred type of an
  // `mlir::Value` to the type resolver by invoking the resolver with
  // the owning operation of the value if it's a block argument.
  void doInitializeForValue(mlir::Value v) {
    if (mlir::BlockArgument ba = llvm::dyn_cast<mlir::BlockArgument>(v)) {
      doVisitOperation(ba.getOwner()->getParentOp());
    }
  }

  // Prints an indentation composed of `indent` times `" "`.
  void printIndent(int indent) {
    for (int i = 0; i < indent; i++)
      llvm::dbgs() << "  ";
  }

  // Dumps the state of type inference for the operation `op` with an
  // indentation level of `indent` as the name of the operation,
  // followed by the types inferred for each operand, followed by
  // `->`, followed by a dump of the state for any operation nested in
  // any region of `op`.
  void dumpStateForOp(mlir::Operation *op, int indent) {
    const LocalInferenceState state = getCurrentInferredTypes(op);

    printIndent(indent);
    llvm::dbgs() << op->getName() << " {";

    llvm::interleaveComma(
        op->getAttrs(), llvm::dbgs(), [&](const mlir::NamedAttribute &attr) {
          llvm::dbgs() << attr.getName() << " = " << attr.getValue();
        });

    llvm::dbgs() << "} : (";

    llvm::interleaveComma(op->getOperands(), llvm::dbgs(), [&](mlir::Value v) {
      llvm::dbgs() << state.find(v);
    });

    llvm::dbgs() << ") -> (";

    llvm::interleaveComma(op->getResults(), llvm::dbgs(), [&](mlir::Value v) {
      llvm::dbgs() << state.find(v);
    });

    llvm::dbgs() << ")\n";

    for (mlir::Region &r : op->getRegions())
      for (mlir::Block &b : r.getBlocks())
        for (mlir::Operation &childOp : b.getOperations())
          dumpStateForOp(&childOp, indent + 1);
  }

  // Dumps the entire state of type inference for the function
  // containing the operation `op`. For each operation, this prints
  // the name of the operation, followed by the types inferred for
  // each operand, followed by `->`, followed by the types inferred
  // for the results.
  void dumpAllState(mlir::ModuleOp module) { dumpStateForOp(module, 0); }

  TypeResolver &resolver;
};

// Type inference analysis running forward by following the flow of
// data from sources to sinks. For each operation encountered, the
// type resolver is invoked in order to update types until
// convergence. May be combined with `BackwardTypeInferenceAnalysis`
// for bidirectional type inference.
class ForwardTypeInferenceAnalysis
    : public TypeInferenceAnalysisBase<
          mlir::dataflow::SparseDataFlowAnalysis<TypeInferenceLattice>> {
public:
  ForwardTypeInferenceAnalysis(mlir::DataFlowSolver &solver,
                               TypeResolver &resolver)
      : TypeInferenceAnalysisBase(resolver, solver) {}

  void setToEntryState(TypeInferenceLattice *lattice) override {
    TypeInferenceAnalysisBase::doInitializeForValue(lattice->getPoint());
  }

  void
  visitOperation(Operation *op,
                 ArrayRef<const TypeInferenceLattice *> latticeOperandTypes,
                 ArrayRef<TypeInferenceLattice *> latticeResultTypes) override {
    TypeInferenceAnalysisBase::doVisitOperation(op);
  }
};

// Type inference analysis running backward by tracking back the flow
// of data from sinks to sources. For each operation encountered, the
// type resolver is invoked in order to update types until
// convergence. May be combined with `ForwardTypeInferenceAnalysis`
// for bidirectional type inference.
class BackwardTypeInferenceAnalysis
    : public TypeInferenceAnalysisBase<
          mlir::dataflow::SparseBackwardDataFlowAnalysis<
              TypeInferenceLattice>> {
public:
  BackwardTypeInferenceAnalysis(mlir::DataFlowSolver &solver,
                                mlir::SymbolTableCollection &symTabs,
                                TypeResolver &resolver)
      : TypeInferenceAnalysisBase(resolver, solver, symTabs) {}

  void setToExitState(TypeInferenceLattice *lattice) override {
    TypeInferenceAnalysisBase::doInitializeForValue(lattice->getPoint());
  }

  void visitOperation(
      Operation *op, ArrayRef<TypeInferenceLattice *> latticeOperandTypes,
      ArrayRef<const TypeInferenceLattice *> latticeResultTypes) override {
    TypeInferenceAnalysisBase::doVisitOperation(op);
  }

  void visitBranchOperand(mlir::OpOperand &operand) override {}
};

} // namespace concretelang
} // namespace mlir

#endif
