open Format
open Common
open Pp

type 'a expr = { expr_desc : 'a expr_desc; expr_meta : 'a }

and 'a expr_desc =
  | Ret of 'a cexpr
  | Let of 'a binding * 'a expr
  | Letrec of 'a binding list * 'a expr
  | Seq of 'a expr list

and 'a cexpr = { cexpr_desc : 'a cexpr_desc; cexpr_meta : 'a }

and 'a cexpr_desc =
  | Atom of 'a atom
  | App of 'a atom * 'a atom list
  | Lam of string list * 'a expr
  | Prim1 of Prim.op1 * 'a cexpr
  | Prim2 of Prim.op2 * 'a atom * 'a atom
  | If of 'a atom * 'a expr * 'a expr
  | Match of 'a atom * 'a pattern list * Type.t
  | Tuple of 'a atom list
  | Get of 'a cexpr * int
  | Vect of 'a atom * 'a atom
  | VectGet of 'a atom * 'a atom
  | VectSet of 'a atom * 'a atom * 'a atom
  | VectLength of 'a cexpr
  | Ascribe of 'a cexpr * Type.t
  | AscribeDyn of 'a cexpr * 'a atom * Inversion.t
  | For of string * 'a atom * 'a atom * 'a expr
  | Variant of string * 'a atom list * Type.t

and 'a atom = { atom_desc : 'a atom_desc; atom_meta : 'a }
and 'a atom_desc = Const of Const.t | Var of string
and 'a binding = { name : string; cexpr : 'a cexpr }
and 'a pattern = { casename : string; args : string list; body : 'a expr }

let expr_of_desc expr_desc expr_meta = { expr_desc; expr_meta }
let atom_of_desc atom_desc atom_meta = { atom_desc; atom_meta }

let cexpr_of_desc cexpr_desc (cexpr_meta : 'a) : 'a cexpr =
  { cexpr_desc; cexpr_meta }

let cexpr_of_atom_desc atom_desc (cexpr_meta : 'a) : 'a cexpr =
  { cexpr_desc = Atom (atom_of_desc atom_desc cexpr_meta); cexpr_meta }

let bind binding expr : 'a expr =
  { expr_desc = Let (binding, expr); expr_meta = expr.expr_meta }

let bind_rec bindings expr : 'a expr =
  { expr_desc = Letrec (bindings, expr); expr_meta = expr.expr_meta }

let ret cexpr (expr_meta : 'a) : 'a expr = { expr_desc = Ret cexpr; expr_meta }

let rec anf (t : 'a Ast.t) : 'a expr =
  match t.desc with
  | Ast.Const const -> ret (cexpr_of_atom_desc (Const const) t.meta) t.meta
  | Ast.Var x -> ret (cexpr_of_atom_desc (Var x) t.meta) t.meta
  | Ast.Lam (args, body) ->
      ret (cexpr_of_desc (Lam (args, anf body)) t.meta) t.meta
  | Ast.App (t1, terms) ->
      anf_atom t1 (fun t1 ->
          anf_atom_list terms (fun terms ->
              ret (cexpr_of_desc (App (t1, terms)) t.meta) t.meta))
  | Ast.Prim1 (op, t1) ->
      anf_cexpr t1 (fun t1 ->
          ret (cexpr_of_desc (Prim1 (op, t1)) t.meta) t.meta)
  | Ast.Prim2 (op, t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              ret (cexpr_of_desc (Prim2 (op, t1, t2)) t.meta) t.meta))
  | Ast.If (t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          ret (cexpr_of_desc (If (t1, anf t2, anf t3)) t.meta) t.meta)
  | Ast.Let (binding, body) ->
      anf_cexpr binding.term (fun term ->
          expr_of_desc
            (Let ({ name = binding.name; cexpr = term }, anf body))
            t.meta)
  | Ast.Letrec (bindings, body) ->
      let open Ast in
      let rec go acc bindings =
        match bindings with
        | [] -> expr_of_desc (Letrec (List.rev acc, anf body)) t.meta
        | binding :: tl ->
            anf_cexpr binding.term (fun term ->
                let binding' = { name = binding.name; cexpr = term } in
                go (binding' :: acc) tl)
      in
      go [] bindings
  | Ast.Match (t1, patterns, ty) ->
      let patterns =
        List.map
          (fun Ast.{ casename; args; body } ->
            { casename; args; body = anf body })
          patterns
      in
      anf_atom t1 (fun t1 ->
          ret (cexpr_of_desc (Match (t1, patterns, ty)) t.meta) t.meta)
  | Ast.Tuple ts ->
      anf_atom_list ts (fun atoms ->
          ret (cexpr_of_desc (Tuple atoms) t.meta) t.meta)
  | Ast.Get (t1, idx) ->
      anf_cexpr t1 (fun t1 -> ret (cexpr_of_desc (Get (t1, idx)) t.meta) t.meta)
  | Ast.Vect (t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              ret (cexpr_of_desc (Vect (t1, t2)) t.meta) t.meta))
  | Ast.VectGet (t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              ret (cexpr_of_desc (VectGet (t1, t2)) t.meta) t.meta))
  | Ast.VectSet (t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              anf_atom t3 (fun t3 ->
                  ret (cexpr_of_desc (VectSet (t1, t2, t3)) t.meta) t.meta)))
  | Ast.VectLength t1 ->
      anf_cexpr t1 (fun t1 -> ret (cexpr_of_desc (VectLength t1) t.meta) t.meta)
  | Ast.Ascribe (t1, ty) ->
      anf_cexpr t1 (fun t1 ->
          ret (cexpr_of_desc (Ascribe (t1, ty)) t.meta) t.meta)
  | Ast.For (x, t1, t2, body) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              ret (cexpr_of_desc (For (x, t1, t2, anf body)) t.meta) t.meta))
  | Ast.Variant (name, terms, ty) ->
      anf_atom_list terms (fun terms ->
          ret (cexpr_of_desc (Variant (name, terms, ty)) t.meta) t.meta)
  | Ast.Seq terms ->
      { expr_desc = Seq (List.map anf terms); expr_meta = t.meta }

and anf_atom (t : 'a Ast.t) (k : 'a atom -> 'a expr) : 'a expr =
  match t.desc with
  | Ast.Const const -> k (atom_of_desc (Const const) t.meta)
  | Ast.Var x -> k (atom_of_desc (Var x) t.meta)
  | Ast.Lam (args, body) ->
      let name = Gensym.fresh "lam" in
      let cexpr = cexpr_of_desc (Lam (args, anf body)) t.meta in
      let binding = { name; cexpr } in
      let atom = atom_of_desc (Var name) t.meta in
      bind binding (k atom)
  | Ast.App (t1, terms) ->
      anf_atom t1 (fun t1 ->
          anf_atom_list terms (fun terms ->
              let name = Gensym.fresh "app" in
              let cexpr = cexpr_of_desc (App (t1, terms)) t.meta in
              let binding = { name; cexpr } in
              let atom = atom_of_desc (Var name) t.meta in
              bind binding (k atom)))
  | Ast.Prim1 (op, t1) ->
      anf_cexpr t1 (fun t1 ->
          let name = Gensym.fresh "prim1" in
          let cexpr = cexpr_of_desc (Prim1 (op, t1)) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Prim2 (op, t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              let name = Gensym.fresh "prim2" in
              let cexpr = cexpr_of_desc (Prim2 (op, t1, t2)) t.meta in
              let binding = { name; cexpr } in
              let atom = atom_of_desc (Var name) t.meta in
              bind binding (k atom)))
  | Ast.If (t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          let name = Gensym.fresh "if" in
          let cexpr = cexpr_of_desc (If (t1, anf t2, anf t3)) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Let (binding, body) ->
      anf_cexpr binding.term (fun term ->
          let name = binding.name in
          let binding = { name; cexpr = term } in
          bind binding (anf_atom body (fun body -> k body)))
  | Ast.Letrec (bindings, body) ->
      let open Ast in
      let rec go acc bindings =
        match bindings with
        | [] -> bind_rec (List.rev acc) (anf_atom body (fun body -> k body))
        | binding :: tl ->
            anf_cexpr binding.term (fun term ->
                let name = binding.name in
                let binding = { name; cexpr = term } in
                go (binding :: acc) tl)
      in
      go [] bindings
  | Ast.Match (t1, patterns, ty) ->
      let patterns =
        List.map
          (fun Ast.{ casename; args; body } ->
            { casename; args; body = anf body })
          patterns
      in
      anf_atom t1 (fun t1 ->
          let name = Gensym.fresh "match" in
          let cexpr = cexpr_of_desc (Match (t1, patterns, ty)) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Tuple ts ->
      anf_atom_list ts (fun atoms ->
          let name = Gensym.fresh "tuple" in
          let cexpr = cexpr_of_desc (Tuple atoms) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Get (t1, idx) ->
      anf_cexpr t1 (fun t1 ->
          let name = Gensym.fresh "get" in
          let cexpr = cexpr_of_desc (Get (t1, idx)) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Vect (t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              let name = Gensym.fresh "vect" in
              let cexpr = cexpr_of_desc (Vect (t1, t2)) t.meta in
              let binding = { name; cexpr } in
              let atom = atom_of_desc (Var name) t.meta in
              bind binding (k atom)))
  | Ast.VectGet (t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              let name = Gensym.fresh "vect_get" in
              let cexpr = cexpr_of_desc (VectGet (t1, t2)) t.meta in
              let binding = { name; cexpr } in
              let atom = atom_of_desc (Var name) t.meta in
              bind binding (k atom)))
  | Ast.VectSet (t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              anf_atom t3 (fun t3 ->
                  let name = Gensym.fresh "vect_get" in
                  let cexpr = cexpr_of_desc (VectSet (t1, t2, t3)) t.meta in
                  let binding = { name; cexpr } in
                  let atom = atom_of_desc (Var name) t.meta in
                  bind binding (k atom))))
  | Ast.VectLength t1 ->
      anf_cexpr t1 (fun t1 ->
          let name = Gensym.fresh "vect_length" in
          let cexpr = cexpr_of_desc (VectLength t1) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Ascribe (t1, ty) ->
      anf_cexpr t1 (fun t1 ->
          let name = Gensym.fresh "ascribed" in
          let cexpr = cexpr_of_desc (Ascribe (t1, ty)) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.For (x, t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              let name = Gensym.fresh "for" in
              let cexpr = cexpr_of_desc (For (x, t1, t2, anf t3)) t.meta in
              let binding = { name; cexpr } in
              let atom = atom_of_desc (Var name) t.meta in
              bind binding (k atom)))
  | Ast.Variant (case_name, terms, ty) ->
      anf_atom_list terms (fun terms ->
          let name = Gensym.fresh "variant" in
          let cexpr = cexpr_of_desc (Variant (case_name, terms, ty)) t.meta in
          let binding = { name; cexpr } in
          let atom = atom_of_desc (Var name) t.meta in
          bind binding (k atom))
  | Ast.Seq terms ->
      let rec go acc terms =
        match terms with
        | [] -> assert false
        | [ term ] ->
            let terms = anf_atom term (fun term -> k term) :: acc |> List.rev in
            { expr_desc = Seq terms; expr_meta = t.meta }
        | term :: tl -> go (anf term :: acc) tl
      in
      if List.length terms > 1 then go [] terms else assert false

and anf_cexpr (t : 'a Ast.t) (k : 'a cexpr -> 'a expr) : 'a expr =
  match t.desc with
  | Ast.Const const -> k (cexpr_of_atom_desc (Const const) t.meta)
  | Ast.Var x -> k (cexpr_of_atom_desc (Var x) t.meta)
  | Ast.Lam (args, body) -> k (cexpr_of_desc (Lam (args, anf body)) t.meta)
  | Ast.App (t1, terms) ->
      anf_atom t1 (fun t1 ->
          anf_atom_list terms (fun terms ->
              k (cexpr_of_desc (App (t1, terms)) t.meta)))
  | Ast.Prim1 (op, t1) ->
      anf_cexpr t1 (fun t1 -> k (cexpr_of_desc (Prim1 (op, t1)) t.meta))
  | Ast.Prim2 (op, t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 -> k (cexpr_of_desc (Prim2 (op, t1, t2)) t.meta)))
  | Ast.If (t1, t2, t3) ->
      anf_atom t1 (fun t1 -> k (cexpr_of_desc (If (t1, anf t2, anf t3)) t.meta))
  | Ast.Let (binding, body) ->
      anf_cexpr binding.term (fun term ->
          let name = binding.name in
          let binding = { name; cexpr = term } in
          bind binding (anf_cexpr body (fun body -> k body)))
  | Ast.Letrec (bindings, body) ->
      let open Ast in
      let rec go acc bindings =
        match bindings with
        | [] -> bind_rec (List.rev acc) (anf_cexpr body (fun body -> k body))
        | binding :: tl ->
            anf_cexpr binding.term (fun term ->
                let name = binding.name in
                let binding = { name; cexpr = term } in
                go (binding :: acc) tl)
      in
      go [] bindings
  | Ast.Match (t1, patterns, ty) ->
      let patterns =
        List.map
          (fun Ast.{ casename; args; body } ->
            { casename; args; body = anf body })
          patterns
      in
      anf_atom t1 (fun t1 ->
          k (cexpr_of_desc (Match (t1, patterns, ty)) t.meta))
  | Ast.Tuple ts ->
      anf_atom_list ts (fun ts -> k (cexpr_of_desc (Tuple ts) t.meta))
  | Ast.Get (t1, idx) ->
      anf_cexpr t1 (fun t1 -> k (cexpr_of_desc (Get (t1, idx)) t.meta))
  | Ast.Vect (t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 -> k (cexpr_of_desc (Vect (t1, t2)) t.meta)))
  | Ast.VectGet (t1, t2) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 -> k (cexpr_of_desc (VectGet (t1, t2)) t.meta)))
  | Ast.VectSet (t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              anf_atom t3 (fun t3 ->
                  k (cexpr_of_desc (VectSet (t1, t2, t3)) t.meta))))
  | Ast.VectLength t1 ->
      anf_cexpr t1 (fun t1 -> k (cexpr_of_desc (VectLength t1) t.meta))
  | Ast.Ascribe (t1, ty) ->
      anf_cexpr t1 (fun t1 -> k (cexpr_of_desc (Ascribe (t1, ty)) t.meta))
  | Ast.For (x, t1, t2, t3) ->
      anf_atom t1 (fun t1 ->
          anf_atom t2 (fun t2 ->
              k (cexpr_of_desc (For (x, t1, t2, anf t3)) t.meta)))
  | Ast.Variant (case_name, terms, ty) ->
      anf_atom_list terms (fun terms ->
          k (cexpr_of_desc (Variant (case_name, terms, ty)) t.meta))
  | Ast.Seq terms ->
      let rec go acc terms =
        match terms with
        | [] -> assert false
        | [ term ] ->
            let terms =
              anf_cexpr term (fun term -> k term) :: acc |> List.rev
            in
            { expr_desc = Seq terms; expr_meta = t.meta }
        | term :: tl -> go (anf term :: acc) tl
      in
      if List.length terms > 1 then go [] terms else assert false

and anf_atom_list (ts : 'a Ast.t list) (k : 'a atom list -> 'a expr) : 'a expr =
  let rec go acc ts =
    match ts with
    | [] -> k (List.rev acc)
    | term :: tl -> anf_atom term (fun term -> go (term :: acc) tl)
  in
  go [] ts

let of_ast (t : 'a Ast.t) : 'a expr =
  Gensym.init ();
  anf t

module StringSet = Set.Make (String)

let rec free_vars e = StringSet.elements (free_vars_e e)
and free_vars_cexpr c = StringSet.elements (free_vars_c c)
and free_vars_atom a = StringSet.elements (free_vars_a a)

and free_vars_e e =
  match e.expr_desc with
  | Ret c -> free_vars_c c
  | Let (binding, e) ->
      StringSet.union
        (free_vars_c binding.cexpr)
        (StringSet.remove binding.name (free_vars_e e))
  | Letrec (bindings, e) ->
      let bound_variables =
        StringSet.of_list (List.map (fun b -> b.name) bindings)
      in
      let exprs = List.map (fun b -> free_vars_c b.cexpr) bindings in
      let used_variables : StringSet.t =
        List.fold_left StringSet.union (free_vars_e e) exprs
      in
      StringSet.diff used_variables bound_variables
  | Seq exprs ->
      List.fold_left
        (fun acc e -> StringSet.union acc (free_vars_e e))
        StringSet.empty exprs

and free_vars_c c =
  match c.cexpr_desc with
  | AscribeDyn (c, a, _) -> StringSet.union (free_vars_c c) (free_vars_a a)
  | Atom a -> free_vars_a a
  | Prim1 (_, c) | Get (c, _) | Ascribe (c, _) | VectLength c -> free_vars_c c
  | App (a1, atoms) ->
      List.fold_left StringSet.union (free_vars_a a1)
        (List.map free_vars_a atoms)
  | Vect (a1, a2) | VectGet (a1, a2) | Prim2 (_, a1, a2) ->
      StringSet.union (free_vars_a a1) (free_vars_a a2)
  | Lam (args, e) -> StringSet.diff (free_vars_e e) (StringSet.of_list args)
  | VectSet (a1, a2, a3) ->
      StringSet.union (free_vars_a a1)
        (StringSet.union (free_vars_a a2) (free_vars_a a3))
  | Tuple atoms | Variant (_, atoms, _) ->
      List.fold_left StringSet.union StringSet.empty
        (List.map free_vars_a atoms)
  | Match (a, patterns, _) ->
      patterns
      |> List.map (fun { args; body; _ } ->
             StringSet.diff (free_vars_e body) (StringSet.of_list args))
      |> List.fold_left StringSet.union (free_vars_a a)
  | If (c, e1, e2) ->
      StringSet.union (free_vars_a c)
        (StringSet.union (free_vars_e e1) (free_vars_e e2))
  | For (id, a1, a2, e) ->
      StringSet.union (free_vars_a a1)
        (StringSet.union (free_vars_a a2) (StringSet.remove id (free_vars_e e)))

and free_vars_a a =
  match a.atom_desc with
  | Const _ -> StringSet.empty
  | Var x -> if x <> "self" then StringSet.singleton x else StringSet.empty

let pp_with_meta ~show_meta pp_meta pp_desc get_meta fmt x =
  if show_meta then fprintf fmt "@[<hv>%a@ %a@]" pp_meta (get_meta x) pp_desc x
  else pp_desc fmt x

let rec pp_expr ?(show_meta = false) pp_meta fmt e =
  pp_with_meta ~show_meta pp_meta
    (pp_expr_desc ~show_meta pp_meta)
    (fun e -> e.expr_meta)
    fmt e

and pp_expr_desc ?(show_meta = false) pp_meta fmt e =
  match e.expr_desc with
  | Ret cexpr -> (pp_cexpr ~show_meta pp_meta) fmt cexpr
  | Let (binding, expr) ->
      fprintf fmt "@[<v>@[<b>let@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_binding ~show_meta pp_meta)
        binding
        (pp_expr ~show_meta pp_meta)
        expr
  | Letrec (bindings, expr) ->
      fprintf fmt "@[<v>@[<b>let rec@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_list ~sep:"@;<1 0>and " (pp_binding ~show_meta pp_meta))
        bindings
        (pp_expr ~show_meta pp_meta)
        expr
  | Seq exprs ->
      fprintf fmt "@[<v>(%a)@]"
        (pp_list ~sep:";@;<1 1>" (fun fmt (e : 'a) ->
             fprintf fmt "@[%a@]" (pp_expr ~show_meta pp_meta) e))
        exprs

and pp_cexpr ?(show_meta = false) pp_meta fmt c =
  pp_with_meta ~show_meta pp_meta
    (pp_cexpr_desc ~show_meta pp_meta)
    (fun c -> c.cexpr_meta)
    fmt c

and pp_cexpr_desc ?(show_meta = false) pp_meta fmt c =
  let pp = pp_cexpr ~show_meta pp_meta in
  match c.cexpr_desc with
  | Atom atom -> pp_atom ~show_meta pp_meta fmt atom
  | App (atom1, atoms) ->
      pp_app (pp_atom ~show_meta pp_meta) fmt (atom1 :: atoms)
  | Lam (args, expr) ->
      fprintf fmt "@[(fun %a ->@;<1 2>%a)@]"
        (pp_list ~sep:" " pp_print_string)
        args
        (pp_expr ~show_meta pp_meta)
        expr
  | Prim1 (op, cexpr) -> fprintf fmt "%a %a" Prim.pp_op1 op pp cexpr
  | Prim2 (op, atom1, atom2) ->
      fprintf fmt "@[(%a %a@;<1 2>%a)@]"
        (pp_atom ~show_meta pp_meta)
        atom1 Prim.pp_op2 op
        (pp_atom ~show_meta pp_meta)
        atom2
  | If (atom, expr1, expr2) ->
      fprintf fmt "@[<hv>(if %a then@;<1 2>%a@;<1 1>else@;<1 2>%a)@]"
        (pp_atom ~show_meta pp_meta)
        atom
        (pp_expr ~show_meta pp_meta)
        expr1
        (pp_expr ~show_meta pp_meta)
        expr2
  | Match (atom, patterns, _) ->
      fprintf fmt "@[<hv>@[match %a@ with@]@;<1 0>| %a@]"
        (pp_atom ~show_meta pp_meta)
        atom
        (pp_list ~sep:"@;<1 0>| " (pp_pattern ~show_meta pp_meta))
        patterns
  | Tuple atoms ->
      fprintf fmt "@[(%a)@]"
        (pp_list ~sep:",@;<1 2>" (pp_atom ~show_meta pp_meta))
        atoms
  | Get (cexpr, i) -> fprintf fmt "%a.(%d)" pp cexpr i
  | Vect (atom1, atom2) ->
      pp_app_name "vector" (pp_atom ~show_meta pp_meta) fmt [ atom1; atom2 ]
  | VectGet (atom1, atom2) ->
      pp_app_name "vector_get" (pp_atom ~show_meta pp_meta) fmt [ atom1; atom2 ]
  | VectSet (atom1, atom2, atom3) ->
      pp_app_name "vector_set"
        (pp_atom ~show_meta pp_meta)
        fmt [ atom1; atom2; atom3 ]
  | VectLength cexpr -> pp_app_name "vector_length" pp fmt [ cexpr ]
  | Ascribe (cexpr, t) ->
      fprintf fmt "@[<hv>(%a@;<0 1>%a@;<1 1>:: %a)@]" pp_meta c.cexpr_meta pp
        cexpr Type.pp t
  | AscribeDyn (cexpr, atom, inv) ->
      fprintf fmt "@[<hv>(%a@;<1 1>:: %a(%a))@]" pp cexpr Inversion.pp inv
        (pp_atom ~show_meta pp_meta)
        atom
  | For (x, atom1, atom2, expr) ->
      fprintf fmt "@[<v>@[for %s = %a to %a do@]@;<1 2>%a@;<1 0>done@]" x
        (pp_atom ~show_meta pp_meta)
        atom1
        (pp_atom ~show_meta pp_meta)
        atom2
        (pp_expr ~show_meta pp_meta)
        expr
  | Variant (name, atoms, _) ->
      if List.length atoms = 0 then fprintf fmt "%s" name
      else
        fprintf fmt "@[%s (@[%a@])@]" name
          (pp_list ~sep:",@ " (pp_atom ~show_meta pp_meta))
          atoms

and pp_atom ?(show_meta = false) pp_meta fmt a =
  pp_with_meta ~show_meta pp_meta pp_atom_desc (fun a -> a.atom_meta) fmt a

and pp_atom_desc fmt a =
  match a.atom_desc with
  | Const const -> fprintf fmt "%s" (Const.to_string const)
  | Var x -> fprintf fmt "%s" x

and pp_binding ?(show_meta = false) pp_meta fmt binding =
  fprintf fmt "%s =@;<1 2>%a" binding.name
    (pp_cexpr ~show_meta pp_meta)
    binding.cexpr

and pp_pattern ?(show_meta = false) pp_meta fmt pattern =
  fprintf fmt "@[%s %a ->@ %a@]" pattern.casename
    (pp_list ~sep:" " pp_print_string)
    pattern.args
    (pp_expr ~show_meta pp_meta)
    pattern.body
