(** This module adds runtime functions in a [Clight] program. These functions
    implement unsupported functions by the target architecture that introduce a
    branch. We need to define them at the [Clight] level in order to have a
    correct labelling. *)


let error_prefix = "Adding runtime functions"
let error = Error.global_error error_prefix

let add_program p s =
  let output = s ^ (ClightPrinter.print_program p) in
  ClightParser.process (`Source ("add_program",  output)) 

type operation =
  | Unary of Clight.unary_operation * Clight.ctype
  | Binary of Clight.binary_operation * Clight.ctype * Clight.ctype
  | Cast of Clight.ctype (* destination type *) * Clight.ctype (* source type *)
  | Fun of string (* name of the function *)

type op_replacement =
    (* operation to be replaced *)
    operation *
    (* base name of the replacement function *)
    string *
    (* C definition of the replacement function, provided a name for the
       function *)
    (string -> string) *
    (* dependencies *)
    operation list

let string_of_operation = function
  | Unary (unop, ctype) ->
    (ClightPrinter.string_of_unop unop) ^ "(" ^
      (ClightPrinter.string_of_ctype ctype) ^ ")"
  | Binary (binop, ctype1, ctype2) ->
    (ClightPrinter.string_of_binop binop) ^ "(" ^
      (ClightPrinter.string_of_ctype ctype1) ^ ", " ^
      (ClightPrinter.string_of_ctype ctype1) ^ ")"
  | Cast (ctype1, ctype2) ->
    "Cast " ^ (ClightPrinter.string_of_ctype ctype1) ^ " " ^
      (ClightPrinter.string_of_ctype ctype2)
  | Fun fun_name -> "Fun " ^ fun_name

let string_of_op_replacement (replaced, base_name, c_def, _) =
  Printf.sprintf "Replaced: %s\n%s\n"
    (string_of_operation replaced) (c_def base_name)

module CtypeSet =
  Set.Make (struct type t = Clight.ctype let compare = Pervasives.compare end)


let deps op replacements =
  let f res (op', _, _, deps) = if op' = op then deps else res in
  List.fold_left f [] replacements


(* Filter used operations only *)

module OperationSet =
  Set.Make (struct type t = operation let compare = Pervasives.compare end)

let union_list l = List.fold_left OperationSet.union OperationSet.empty l

let f_ctype ctype _ = ctype

let f_expr e _ sub_expr_descrs_res =
  let res_e = match e with
    | Clight.Expr (Clight.Evar x, Clight.Tfunction _) ->
      OperationSet.singleton (Fun x)
    | _ -> OperationSet.empty in
  union_list (res_e :: sub_expr_descrs_res)

let f_expr_descr ed _ sub_exprs_res =
  let res_ed = match ed with
    | Clight.Eunop (unop, Clight.Expr (_, t)) ->
      OperationSet.singleton (Unary (unop, t))
    | Clight.Ebinop (binop, Clight.Expr (_, t1), Clight.Expr (_, t2)) ->
      OperationSet.singleton (Binary (binop, t1, t2))
    | Clight.Ecast (t, Clight.Expr (_, t')) ->
      OperationSet.singleton (Cast (t, t'))
    | _ -> OperationSet.empty in
  union_list (res_ed :: sub_exprs_res)

let f_stmt _ sub_exprs_res sub_stmts_res =
  OperationSet.union (union_list sub_exprs_res) (union_list sub_stmts_res)

let used_ops_stmt = ClightFold.statement f_ctype f_expr f_expr_descr f_stmt

let used_ops_fundef = function
  | Clight.Internal cfun -> used_ops_stmt cfun.Clight.fn_body
  | Clight.External _ -> OperationSet.empty

let used_ops p =
  let f ops (_, fundef) = OperationSet.union ops (used_ops_fundef fundef) in
  List.fold_left f OperationSet.empty p.Clight.prog_funct


let mem_replacements op =
  let f res (op', _, _, _) = res || (op' = op) in
  List.fold_left f false

let rec fix equal next x =
  let y = next x in
  if equal x y then x
  else fix equal next y

let needed_replacements used_ops replacements =
  let f op = mem_replacements op replacements in
  let reduced_used_ops = OperationSet.filter f used_ops in
  let next_ops ops =
    let add ops op = OperationSet.add op ops in
    let f op next_ops = List.fold_left add next_ops (deps op replacements) in
    OperationSet.fold f ops ops in
  let needed_ops = fix OperationSet.equal next_ops reduced_used_ops in
  let f (op, _, _, _) = OperationSet.mem op needed_ops in
  List.filter f replacements


let map_fresh_name unique map base_name =
  if StringTools.Map.mem base_name map then
    (map, StringTools.Map.find base_name map)
  else
    let fresh_name = unique base_name in
    (StringTools.Map.add base_name fresh_name map, fresh_name)

let fresh_replacements unique replacements =
  let f (map, l) (op, base_name, new_fun, deps) =
    let (map, fresh_name) = map_fresh_name unique map base_name in
    (map, l @ [(op, fresh_name, new_fun fresh_name, deps)]) in
  snd (List.fold_left f (StringTools.Map.empty, []) replacements)


let add_replacements replacements p =
  let f_replacements s (_, _, new_fun, _) = s ^ "\n" ^ new_fun in
  let added_string = List.fold_left f_replacements "" replacements in
  add_program p added_string


let make_replacement_assoc = List.map (fun (op, name, _, _) -> (op, name))


let add p replacements =
  let used_ops = used_ops p in
  let needed_replacements = needed_replacements used_ops replacements in
  let unique = StringTools.make_unique (ClightAnnotator.all_idents p) in
  let replacements = fresh_replacements unique needed_replacements in
  let p = add_replacements replacements p in
  (p, make_replacement_assoc replacements)


(* Replacement *)

let seq =
  List.fold_left
    (fun stmt1 stmt2 ->
      if stmt1 = Clight.Sskip then stmt2
      else
	if stmt2 = Clight.Sskip then stmt1
	else Clight.Ssequence (stmt1, stmt2))
    Clight.Sskip

let f_ctype ctype _ = ctype

let call fresh replaced replacement_assoc args ret_type =
  let tmp = fresh () in
  let replacement_fun_name = List.assoc replaced replacement_assoc in
  let tmpe = Clight.Expr (Clight.Evar tmp, ret_type) in
  let (args, args_types) = List.split args in
  let f_type = Clight.Tfunction (args_types, ret_type) in
  let f = Clight.Expr (Clight.Evar replacement_fun_name, f_type) in
  let call = Clight.Scall (Some tmpe, f, args) in
  (tmpe, call, [(tmp, ret_type)])

let replace_ident replacement_assoc x t =
  let new_name = match t with
    | Clight.Tfunction _
	when List.mem_assoc (Fun x) replacement_assoc ->
      let replacement_fun_name = List.assoc (Fun x) replacement_assoc in
      replacement_fun_name
    | _ -> x in
  (Clight.Expr (Clight.Evar new_name, t), Clight.Sskip, [])

let f_expr fresh replacement_assoc e _ _ =

  let f (Clight.Expr (ed, t) as e) sub_exprs_res =
    let f_sub_exprs (es, stmt, tmps) (e, stmt', tmps') =
      (es @ [e], seq [stmt ; stmt'], tmps @ tmps') in
    let (sub_exprs, stmt1, tmps1) =
      List.fold_left f_sub_exprs ([], Clight.Sskip, []) sub_exprs_res in
    let (e, stmt2, tmps2) = match ed, sub_exprs with
      | Clight.Evar x, _ -> replace_ident replacement_assoc x t
      | Clight.Eunop (unop, Clight.Expr (_, t')), e' :: _
	  when List.mem_assoc (Unary (unop, t')) replacement_assoc ->
	call fresh (Unary (unop, t')) replacement_assoc [(e', t')] t
      | Clight.Ebinop (binop, Clight.Expr (_, t1), Clight.Expr (_, t2)),
	e1 :: e2 :: _
	  when List.mem_assoc (Binary (binop, t1, t2)) replacement_assoc ->
	call fresh (Binary (binop, t1, t2)) replacement_assoc
	  [(e1, t1) ; (e2, t2)] t
      | Clight.Ecast (t, Clight.Expr (_, t')), e' :: _
	  when List.mem_assoc (Cast (t, t')) replacement_assoc ->
	call fresh (Cast (t, t')) replacement_assoc [(e', t')] t
      | _ -> (e, Clight.Sskip, []) in
    (ClightFold.expr_fill_exprs e sub_exprs,
     seq [stmt1 ; stmt2],
     tmps1 @ tmps2) in

  ClightFold.expr2 f e

let f_expr_descr ed _ _ _ = ed

let f_stmt stmt sub_exprs_res sub_stmts_res =
  let f_sub_exprs (es, stmt, tmps) (e, stmt', tmps') =
    (es @ [e], seq [stmt ; stmt'], tmps @ tmps') in
  let (sub_exprs, added_stmt, tmps1) =
    List.fold_left f_sub_exprs ([], Clight.Sskip, []) sub_exprs_res in
  let f_sub_stmts (stmts, tmps) (stmt, tmps') =
    (stmts @ [stmt], tmps @ tmps') in
  let (sub_stmts, tmps2) = List.fold_left f_sub_stmts ([], []) sub_stmts_res in
  let stmt' = ClightFold.statement_fill_subs stmt sub_exprs sub_stmts in
  (seq [added_stmt ; stmt'], tmps1 @ tmps2)

let replace_statement fresh replacement_assoc =
  ClightFold.statement f_ctype (f_expr fresh replacement_assoc)
    f_expr_descr f_stmt

let replace_internal fresh replacement_assoc def =
  let (new_body, tmps) =
    replace_statement fresh replacement_assoc def.Clight.fn_body in
  { def with
    Clight.fn_vars = def.Clight.fn_vars @ tmps ; Clight.fn_body = new_body }

let replace_funct fresh replacement_assoc (id, fundef) =
  let fundef' = match fundef with
    | Clight.Internal def ->
      Clight.Internal (replace_internal fresh replacement_assoc def)
    | _ -> fundef in
  (id, fundef')

let replace fresh replacement_assoc p =
  let prog_funct =
    List.map (replace_funct fresh replacement_assoc) p.Clight.prog_funct in
  { p with Clight.prog_funct = prog_funct }


let save_and_parse p =
  try
    ClightParser.process (`Source ("replaced", ClightPrinter.print_program p))
  with Sys_error _ -> failwith "Error reparsing Clight8 transformation."

let add_replacements p replacements =
  let p = ClightCasts.simplify p in
  let (p, replacement_assoc) = add p replacements in
  let p = ClightCasts.simplify p in
  let tmp_universe = ClightAnnotator.fresh_universe "_tmp" p in
  let fresh () = StringTools.Gen.fresh tmp_universe in
  let p = replace fresh replacement_assoc p in
  let p = save_and_parse p in
  ClightCasts.simplify p


(* Unsupported operations by the 8051. *)

let cint size sign = Clight.Tint (size, sign)

let cints size = cint size AST.Signed
let cintu size = cint size AST.Unsigned

let cint8s = cints Clight.I8
let cint8u = cintu Clight.I8
let cint16s = cints Clight.I16
let cint16u = cintu Clight.I16
let cint32s = cints Clight.I32
let cint32u = cintu Clight.I32

let byte_size_of_intsize = function
  | Clight.I8 -> 1
  | Clight.I16 -> 2
  | Clight.I32 -> 4

let bit_size_of_intsize size = (byte_size_of_intsize size) * 8

let string_of_intsize size = string_of_int (bit_size_of_intsize size)

let ctype_of_intsize = function
  | Clight.I8 -> "char"
  | Clight.I16 -> "short"
  | Clight.I32 -> "int"


(* Unsigned divisions and modulos *)

let divumodu_fun res t s =
  "unsigned " ^ t ^ " " ^ s ^ " (unsigned " ^ t ^ " x, unsigned " ^ t ^ " y)" ^
    "{\n" ^
  "  unsigned " ^ t ^ " quo = 0;\n" ^
  "  unsigned " ^ t ^ " rem = x;\n" ^
  "  while (rem >= y) {\n" ^
  "    quo = quo + 1;\n" ^
  "    rem = rem - y;\n" ^
  "  }\n" ^
  "  return (" ^ res ^ ");\n" ^
  "}\n\n"

let divumodu op sizes sizec t =
  let (prefix, res) = match op with
    | Clight.Odiv -> ("div", "quo")
    | Clight.Omod -> ("mod", "rem")
    | _ -> assert false (* do not use on these arguments *) in
  let cu = cintu sizec in
  (Binary (op, cu, cu), "_" ^ prefix ^ sizes ^ "u", divumodu_fun res t, [])

let divu = divumodu Clight.Odiv
let modu = divumodu Clight.Omod


(* Unsigned divisions *)

(* 16 bits unsigned division *)

let div16u = divu "16" Clight.I16 "short"

(* 32 bits unsigned division *)

let div32u = divu "32" Clight.I32 "int"

(* Signed divisions *)

let divs_fun t s =
  "signed " ^ t ^ " " ^ s ^ " (signed " ^ t ^ " x, signed " ^ t ^ " y) {\n" ^
  "  unsigned " ^ t ^ " x1 = (unsigned " ^ t ^ ") x;\n" ^
  "  unsigned " ^ t ^ " y1 = (unsigned " ^ t ^ ") y;\n" ^
  "  signed " ^ t ^ " sign = 1;\n" ^
  "  if (x < 0) { x1 = (unsigned " ^ t ^ ") (-x); sign = -sign; }\n" ^
  "  if (y < 0) { y1 = (unsigned " ^ t ^ ") (-y); sign = -sign; }\n" ^
  "  unsigned " ^ t ^ " res = (x1/y1);\n" ^
  "  return (sign * ((signed " ^ t ^ ") res));\n" ^
  "}\n\n"

let divs sizes sizec t =
  let cs = cints sizec in
  let cu = cintu sizec in
  (Binary (Clight.Odiv, cs, cs), "_div" ^ sizes ^ "s", divs_fun t,
   [Binary (Clight.Odiv, cu, cu)])

(* 8 bits signed division *)

let div8s = divs "8" Clight.I8 "char"

(* 16 bits signed division *)

let div16s = divs "16" Clight.I16 "short"

(* 32 bits signed division *)

let div32s = divs "32" Clight.I32 "int"


(* Unsigned modulos *)

(* 16 bits unsigned modulo *)

let mod16u = modu "16" Clight.I16 "short"

(* 32 bits unsigned modulo *)

let mod32u = modu "32" Clight.I32 "int"

(* Signed modulos *)

let mods_fun t s =
  "signed " ^ t ^ " " ^ s ^ " (signed " ^ t ^ " x, signed " ^ t ^ " y) {\n" ^
  "  return (x - (x/y) * y);\n" ^
  "}\n\n"

let mods size ct t =
  (Binary (Clight.Omod, ct, ct), "_mod" ^ size ^ "s", mods_fun t,
   [Binary (Clight.Odiv, ct, ct)])

(* 8 bits signed modulo *)

let mod8s = mods "8" cint8s "char"

(* 16 bits signed modulo *)

let mod16s = mods "16" cint16s "short"

(* 32 bits signed modulo *)

let mod32s = mods "32" cint32s "int"


(* Shifts *)

let sh_fun op t s =
  t ^ " " ^ s ^ " (" ^ t ^ " x, " ^ t ^ " y) {\n" ^
  "  " ^ t ^ " res = x, i;\n" ^
  "  for (i = 0 ; i < y ; i++) res = res " ^ op ^ " 2;\n" ^
  "  return res;\n" ^
  "}\n\n"

let sh cop sop direction deps size sign =
  let sizes = string_of_intsize size in
  let ct = Clight.Tint (size, sign) in
  let (short_sign, long_sign) = match sign with
    | AST.Signed -> ("s", "signed ")
    | AST.Unsigned -> ("u", "unsigned ") in
  let t = long_sign ^ (ctype_of_intsize size) in
  (Binary (cop, ct, ct), "_sh" ^ direction ^ sizes ^ short_sign,
   sh_fun sop t, deps)


(* Shift lefts *)

let shl = sh Clight.Oshl "*" "l" []

(* Signed shift lefts *)

(* 8 bits signed shift left *)

let shl8s = shl Clight.I8 AST.Signed

(* 16 bits signed shift left *)

let shl16s = shl Clight.I16 AST.Signed

(* 32 bits signed shift left *)

let shl32s = shl Clight.I32 AST.Signed

(* Unsigned shift lefts *)

(* 8 bits unsigned shift left *)

let shl8u = shl Clight.I8 AST.Unsigned

(* 16 bits unsigned shift left *)

let shl16u = shl Clight.I16 AST.Unsigned

(* 32 bits unsigned shift left *)

let shl32u = shl Clight.I32 AST.Unsigned


(* Shift rights *)

(* Signed shift rights *)

let shrs_fun size t s =
  "signed " ^ t ^ " " ^ s ^ " (signed " ^ t ^ " x, signed " ^ t ^ " y) {\n" ^
  "  unsigned " ^ t ^ " x1, y1, to_add, res, i;\n" ^
  "  if (y < 0) return 0;\n" ^
  "  x1 = x; y1 = y; to_add = 1; res = x1;" ^
  "  for (i = 0 ; i < " ^ size ^ " ; i++) to_add = to_add * 2;\n" ^
  "  to_add = to_add & x1;\n" ^
  "  for (i = 0 ; i < y1 ; i++) res = (res / 2) + to_add;\n" ^
  "  return ((signed " ^ t ^ ") res);\n" ^
  "}\n\n"

let shrs size =
  let sizes = string_of_int (bit_size_of_intsize size) in
  let op_sizes = string_of_int ((bit_size_of_intsize size) - 1) in
  let t = ctype_of_intsize size in
  let ct = Clight.Tint (size, AST.Signed) in
  let ctdep = Clight.Tint (size, AST.Unsigned) in
  (Binary (Clight.Oshr, ct, ct), "_shr" ^ sizes ^ "s", shrs_fun op_sizes t,
   [Binary (Clight.Odiv, ctdep, ctdep)])

(* 8 bits signed shift right *)

let shr8s = shrs Clight.I8

(* 16 bits signed shift right *)

let shr16s = shrs Clight.I16

(* 32 bits signed shift right *)

let shr32s = shrs Clight.I32

(* Unsigned shift rights *)

let shru size =
  let t_dep = Clight.Tint (size, AST.Unsigned) in
  sh Clight.Oshr "/" "r" [Binary (Clight.Odiv, t_dep, t_dep)] size AST.Unsigned

(* 8 bits unsigned shift right *)

let shr8u = shru Clight.I8

(* 16 bits unsigned shift right *)

let shr16u = shru Clight.I16

(* 32 bits unsigned shift right *)

let shr32u = shru Clight.I32


let unsupported =
  [div16u ; div32u ; div8s ; div16s ; div32s ;
   mod16u ; mod32u ; mod8s ; mod16s ; mod32s ;
   shl8s ; shl16s ; shl32s ; shl8u ; shl16u ; shl32u ;
   shr8s ; shr16s ; shr32s ; shr8u ; shr16u ; shr32u]


let replace_unsupported p = add_replacements p unsupported