open Ast
module Env = Map.Make (String)

let collect_constant_strings code =
  let counter = ref 0 in
  let env = ref Env.empty in
  let rec ccs_expr = function
    | IR.Void -> IR2.Value (IR2.Nil), []
    | IR.Bool b -> IR2.Value (IR2.Bool b), []
    | IR.Int n -> IR2.Value (IR2.Int n), []
    | IR.Char c -> IR2.Value (IR2.Char c), []
    | IR.Float f -> IR2.Value (IR2.Float f), []
    | IR.Str s ->
      (match Env.find_opt s !env with
      | Some l -> IR2.Value (IR2.Data l), []
      | None ->
        incr counter;
        let l = "str" ^ string_of_int !counter in
        env := Env.add s l !env;
        IR2.Value (IR2.Data l), [ l, s ])
    | IR.Var v -> IR2.Var v, []
    | IR.Call (f, args) ->
      let args2 = List.map ccs_expr args in
      let ccs = List.flatten (List.map (fun (_, s) -> s) args2) in
      IR2.Call (f, List.map (fun (e, _) -> e) args2), ccs
  in
  (* let ccs_lvalue = function *)
  (*   | IR.LVar v  -> *)
  (*      IR2.LVar v, [] *)
  (*   | IR.LAddr a -> *)
  (*      let a2, ccs = ccs_expr a in *)
  (*      IR2.LAddr a2, ccs *)
  (* in *)
  let rec ccs_instr = function
    | IR.Decl v -> IR2.Decl v, []
    | IR.Return e ->
      let e2, ccs = ccs_expr e in
      IR2.Return e2, ccs
    | IR.Expr e ->
      let e2, ccs = ccs_expr e in
      IR2.Expr e2, ccs
    | IR.Assign (lv, e) ->
      let lv2, ccs_lv = ccs_expr (IR.Var lv) in
      let e2, ccs_e = ccs_expr e in
      IR2.Assign (lv2, e2), List.flatten [ ccs_lv; ccs_e ]
    | IR.Cond (t, y, n) ->
      let t2, ccs_t = ccs_expr t in
      let y2, ccs_y = ccs_block y in
      let n2, ccs_n = ccs_block n in
      IR2.Cond (t2, y2, n2), List.flatten [ ccs_t; ccs_y; ccs_n ]
    | IR.Loop (e, b) ->
      let e2, ccs_e = ccs_expr e in
      let b2, ccs_b = ccs_block b in
      IR2.Loop (e2, b2), List.flatten [ ccs_e; ccs_b ]
  and ccs_block = function
    | [] -> [], []
    | i :: r ->
      let i2, ccs_i = ccs_instr i in
      let r2, ccs_r = ccs_block r in
      i2 :: r2, List.flatten [ ccs_i; ccs_r ]
  in
  let ccs_def = function
    | IR.Func (name, args, body) ->
      let body2, ccs = ccs_block body in
      IR2.Func (name, args, body2), ccs
  in
  let rec ccs_prog = function
    | [] -> [], []
    | d :: r ->
      let d2, ccs_d = ccs_def d in
      let r2, ccs_r = ccs_prog r in
      d2 :: r2, List.flatten [ ccs_d; ccs_r ]
  in
  ccs_prog code
;;

let simplify code = collect_constant_strings code