open Syntax
open Ir
open Effects

let primop_of_binop = function
  | Add -> PAdd | Sub -> PSub | Mul -> PMul | Div -> PDiv
  | Eq -> PEq | Neq -> PNeq | Lt -> PLt | Le -> PLe
  | Gt -> PGt | Ge -> PGe | And -> PAnd | Or -> POr

let gensym =
  let counter = ref 0 in
  fun () -> incr counter; Printf.sprintf "v_%d" !counter

let rec lower (e : expr) : ir =
  let ir = match e with
  | Var x -> mk (Var x)
  | Int n -> mk (Lit (LInt n))
  | Bool b -> mk (Lit (LBool b))
  | Binop (op, e1, e2) ->
    let args = [lower e1; lower e2] in
    mk (Prim (primop_of_binop op, args))
  | If (c, t, f) ->
    mk (If (lower c, lower t, lower f))
  | Let (x, e1, e2) ->
    mk (Let (x, lower e1, lower e2))
  | LetRec (f, x, body, cont) ->
    mk (Fix ([(f, x, lower body)], lower cont))
  | Fun (x, body) ->
    mk (Lam (x, lower body))
  | App (e1, e2) ->
    mk (App (lower e1, lower e2))
  | Match (scrut, cases) ->
    lower_match (lower scrut) cases
  | Seq (e1, e2) ->
    let x = gensym () in
    mk (Let (x, lower e1, lower e2))
  in
  analyse ir;
  ir

and lower_match scr cases =
  match cases with
  | [] -> mk (Lit (LInt 0))
  | (pat, body) :: rest ->
    let body_ir = lower body in
    let fail = lower_match scr rest in
    match pat with
    | PWildcard -> body_ir
    | PVar x -> mk (Let (x, scr, body_ir))
    | PInt n ->
      let cond = mk (Prim (PEq, [scr; mk (Lit (LInt n))])) in
      analyse cond;
      mk (If (cond, body_ir, fail))
    | PBool b ->
      let cond = mk (Prim (PEq, [scr; mk (Lit (LBool b))])) in
      analyse cond;
      mk (If (cond, body_ir, fail))