open Syntax
open Closure

module type S = sig
  type t

  val codegen_program : t program -> Cmm.Syntax.compunit
end

module type MemRep = sig
  type t
  type ev

  val const : Const.t -> t
  val atom : ev atom -> t
  val allocate_closure : string -> string list -> ev -> t
  val fill_closure : t -> string -> string list -> t
  val closure : string -> string list -> ev -> t
  val bind_free_vars : string list -> t -> t
  val tuple : ev atom list -> t
  val vect : ev atom -> ev atom -> t
  val variant : string -> ev atom list -> Type.t -> t
end

module type Elim = sig
  type t
  type ev

  val prim1 : Prim.op1 -> t -> t
  val prim2 : Prim.op2 -> t -> t -> t
  val get : t -> int -> t
  val ascribe : t -> Type.t -> t
  val call : ?well_known:string option -> t -> t list -> t
  val ite : t -> t -> t -> t
  val match_ : t -> (string * string list * t) list -> Type.t -> t
  val for_ : string -> t -> t -> t -> t
  val vect_get : t -> t -> t
  val vect_set : t -> t -> t -> t
  val vect_length : t -> t
end

module MakeLower
    (MemRep : MemRep with type t = Cmm.Syntax.expr)
    (Elim : Elim with type t := MemRep.t and type ev := MemRep.ev) =
struct
  type t = MemRep.ev

  let rec codegen_program (program : MemRep.ev program) : Cmm.Syntax.compunit =
    let funcs = List.map codegen_func program.funcs in
    let func =
      Cmm.Syntax.
        {
          name = "main";
          args = [];
          return = TInt Int32;
          body = Seq [ codegen_expr program.expr; Int (Int32, 0L) ];
        }
    in
    funcs @ [ func ]

  and codegen_func (func : 'a func) : Cmm.Syntax.fundef =
    match func with
    | Func f ->
        Cmm.Syntax.
          {
            name = f.name;
            args = ("self", TPtr) :: List.map (fun x -> (x, TInt Int64)) f.args;
            return = TInt Int64;
            body = MemRep.bind_free_vars f.free_vars (codegen_expr f.body);
          }

  and codegen_expr (expr : 'a expr) : Cmm.Syntax.expr =
    match expr.expr_desc with
    | Ret cexpr -> codegen_cexpr cexpr
    | Let (binding, expr) ->
        Cmm.Syntax.Let
          (binding.name, codegen_cexpr binding.cexpr, codegen_expr expr)
    | Letrec (bindings, expr) ->
        (* fail if expressions other than lambdas are bound with letrec *)
        let closure_bindings =
          bindings
          |> List.map (fun binding ->
                 match binding.cexpr.cexpr_desc with
                 | Closure (fname, captured_vars) ->
                     ( binding.name,
                       (fname, captured_vars),
                       binding.cexpr.cexpr_meta )
                 | _ ->
                     failwith "expression other than lambda bound with let rec")
        in
        (* allocate bound closures *)
        let rec allocate_closures l ret =
          match l with
          | (id, (fname, captured_vars), ev) :: tl ->
              Cmm.Syntax.Let
                ( id,
                  MemRep.allocate_closure fname captured_vars ev,
                  allocate_closures tl ret )
          | [] -> ret
        in
        (* fill allocated closures *)
        let fill_closures =
          closure_bindings
          |> List.map (fun (id, (fname, captured_vars), _) ->
                 MemRep.fill_closure Cmm.Syntax.(Var id) fname captured_vars)
        in
        allocate_closures closure_bindings
          (Cmm.Syntax.Seq (fill_closures @ [ codegen_expr expr ]))
    | Seq exprs -> Cmm.Syntax.Seq (List.map codegen_expr exprs)

  and codegen_cexpr (cexpr : 'a cexpr) : Cmm.Syntax.expr =
    match cexpr.cexpr_desc with
    | Atom a -> codegen_atom a
    | Call (a1, atoms, well_known) ->
        Elim.call ~well_known (codegen_atom a1) (List.map codegen_atom atoms)
    | Closure (fname, free_vars) ->
        MemRep.closure fname free_vars cexpr.cexpr_meta
    | Prim1 (op1, c1) -> Elim.prim1 op1 (codegen_cexpr c1)
    | Prim2 (op2, a1, a2) -> Elim.prim2 op2 (codegen_atom a1) (codegen_atom a2)
    | If (a, e1, e2) ->
        Elim.ite (codegen_atom a) (codegen_expr e1) (codegen_expr e2)
    | Match (a, patterns, ty) ->
        let a = codegen_atom a in
        let patterns =
          List.map
            (fun Closure.{ casename; args; body } ->
              (casename, args, codegen_expr body))
            patterns
        in
        Elim.match_ a patterns ty
    | Tuple atoms -> MemRep.tuple atoms
    | Get (c, i) -> Elim.get (codegen_cexpr c) i
    | Vect (a1, a2) -> MemRep.vect a1 a2
    | VectGet (a1, a2) -> Elim.vect_get (codegen_atom a1) (codegen_atom a2)
    | VectSet (a1, a2, a3) ->
        Elim.vect_set (codegen_atom a1) (codegen_atom a2) (codegen_atom a3)
    | VectLength c -> Elim.vect_length (codegen_cexpr c)
    | Ascribe (c, ty) -> Elim.ascribe (codegen_cexpr c) ty
    | AscribeDyn (_, _, _) -> failwith "not implemented yet"
    | For (var, start, stop, body) ->
        Elim.for_ var (codegen_atom start) (codegen_atom stop)
          (codegen_expr body)
    | Variant (case_name, atoms, ty) -> MemRep.variant case_name atoms ty

  and codegen_atom (atom : 'a atom) : Cmm.Syntax.expr =
    match atom.atom_desc with
    | Const c -> codegen_const c
    | Var x -> Cmm.Syntax.Var x

  and codegen_const (const : Const.t) : Cmm.Syntax.expr = MemRep.const const
end
