(** MIR types and modules corresponding to the expressions of the language *)

module Pattern : sig
  type litType = Int | Real | Imaginary | Str [@@deriving sexp, hash, compare]

  type 'a t =
    | Var of string
    | Lit of litType * string
    | FunApp of 'a Fun_kind.t * 'a list
    | TernaryIf of 'a * 'a * 'a
    | EAnd of 'a * 'a
    | EOr of 'a * 'a
    | Indexed of 'a * 'a Index.t list
    | Promotion of 'a * UnsizedType.t * UnsizedType.autodifftype
    | TupleProjection of 'a * int
  [@@deriving sexp, hash, compare, map, fold]
end

(** The "two-level" type for statements in the MIR. This corresponds to what the
    AST calls [Frontend.Ast.expr_with] *)
type 'a t = {pattern: 'a t Pattern.t; meta: 'a} [@@deriving compare, hash, sexp]

val pp : 'a t Fmt.t

val rewrite_bottom_up : f:('a t -> 'a t) -> 'a t -> 'a t
(** [rewrite_bottom_up] specializes [fold] so that the result type ['r] is equal
    to the type of our fixed-point data structure i.e. ['r = 'a t]. This also
    means that the function [f] can be written with our fixed-point type ['a t]
    as its argument. *)

module Typed : sig
  module Meta : sig
    type t =
      { type_: UnsizedType.t
      ; loc: Location_span.t [@sexp.opaque] [@compare.ignore]
      ; adlevel: UnsizedType.autodifftype }
    [@@deriving compare, create, sexp, hash]

    val empty : t
  end

  type nonrec t = (Meta.t[@compare.ignore]) t [@@deriving hash, sexp, compare]

  val type_of : t -> UnsizedType.t
  val adlevel_of : t -> UnsizedType.autodifftype
  val fun_arg : t -> UnsizedType.autodifftype * UnsizedType.t
  val pp : t Fmt.t

  include Core.Comparator.S with type t := t

  include
    Core.Comparable.S
      with type t := t
       and type comparator_witness := comparator_witness
end

module Helpers : sig
  val int : int -> Typed.t
  val float : float -> Typed.t
  val complex : float * float -> Typed.t
  val str : string -> Typed.t
  val variable : string -> Typed.t
  val zero : Typed.t
  val one : Typed.t
  val unary_op : Operator.t -> Typed.t -> Typed.t
  val binop : Typed.t -> Operator.t -> Typed.t -> Typed.t
  val binop_list : Typed.t list -> Operator.t -> default:Typed.t -> Typed.t
  val row_vector : float list -> Typed.t
  val vector : float list -> Typed.t
  val matrix : float list list -> Typed.t
  val matrix_from_rows : Typed.t list -> Typed.t
  val complex_row_vector : (float * float) list -> Typed.t
  val complex_vector : (float * float) list -> Typed.t
  val complex_matrix_from_rows : Typed.t list -> Typed.t
  val array_expr : Typed.t list -> Typed.t
  val tuple_expr : Typed.t list -> Typed.t
  val try_unpack : Typed.t -> Typed.t list option
  val loop_bottom : Typed.t
  val internal_funapp : 'a t Internal_fun.t -> 'a t list -> 'a -> 'a t
  val contains_fn_kind : ('a t Fun_kind.t -> bool) -> ?init:bool -> 'a t -> bool
  val infer_type_of_indexed : UnsizedType.t -> 'a Index.t list -> UnsizedType.t
  val add_int_index : Typed.t -> Typed.t Index.t -> Typed.t
  val add_tuple_index : Typed.t -> int -> Typed.t
  val collect_indices : 'a t -> 'a t Index.t list
end
