Source file insert_copy_and_fix_length.ml

open Utils
open Prog
module L = Location

let is_array_copy (x:lval) (e:expr) =
  match x with
  | Lvar x ->
    let x = L.unloc x in
    begin match x.v_ty with
    | Arr (xws, xn) ->
      begin match e with
      | Pvar y ->
        let y = L.unloc y.gv in
        begin match y.v_ty with
        | Arr(yws, yn) ->
           (* Ignore ill-typed copies: they are later rejected by “typing”. *)
           if arr_size yws yn < arr_size xws xn then None else
           if x.v_kind = Reg(Normal, Direct) then Some (xws, xn)
           else if y.v_kind = Reg(Normal, Direct) then Some (yws, arr_size xws xn / size_of_ws yws)
           else None
        | _ -> None
        end
      | _ -> None
      end
    | _ -> None
    end
  | _ -> None

let size_of_lval =
  function
  | Lvar x -> size_of (L.unloc x).v_ty
  | Lasub (_, ws, len, _, _) -> arr_size ws len
  | Lnone _ | Lmem _ | Laset _ -> assert false

let rec fix_length_eassert e =
  match e with
  | Pexpr _ -> e
  | PappN_safety (o, es) ->
    let e = List.hd es in
    let ty = Typing.type_of_expr e in
    let len = Conv.cz_of_int (size_of ty) in
    let o = match o with Ois_arr_init _ -> Operators.Ois_arr_init len | Ois_barr_init _ -> Ois_barr_init len in
    PappN_safety(o, es)
  | Pis_var_init _ | Pis_mem_init _ -> e
  | Pand(e1, e2) -> Pand (fix_length_eassert e1, fix_length_eassert e2)

let rec iac_stmt pd is = List.map (iac_instr pd) is
and iac_instr pd i = { i with i_desc = iac_instr_r pd i.i_loc i.i_desc }
and iac_instr_r pd loc ir =
  match ir with
  | Cassgn (x, t, _, e) ->
    begin match is_array_copy x e with
      | None -> ir
      | Some (ws, n) ->
          warning IntroduceArrayCopy
            loc "an array copy is introduced";
          let op = Pseudo_operator.Ocopy(ws, Conv.cz_of_int n) in
          Copn([x], t, Sopn.Opseudo_op op, [e])
    end
  | Cif (b, th, el) -> Cif (b, iac_stmt pd th, iac_stmt pd el)
  | Cfor (i, r, s) -> Cfor (i, r, iac_stmt pd s)
  | Cwhile (a, c1, t, info, c2) -> Cwhile (a, iac_stmt pd c1, t, info, iac_stmt pd c2)
  | Copn (xs,t,o,es) ->

    begin match o, xs with
    | Sopn.Opseudo_op(Pseudo_operator.Ospill(o,_)), _ ->
      let tys = List.map (fun e -> Conv.cty_of_ty (Typing.ty_expr pd loc e)) es in
      Copn(xs,t, Sopn.Opseudo_op(Pseudo_operator.Ospill(o, tys)), es)

    | Sopn.Opseudo_op(Pseudo_operator.Ocopy(ws, _)), [x] ->
      (* Fix the size it is dummy for the moment *)
      let xn = size_of_lval x in
      let wsn = size_of_ws ws in
      if xn mod wsn <> 0 then
        Typing.error loc
          "the destination %a has size %i: it should be a multiple of %i"
          (Printer.pp_lval ~debug:false) x
          xn wsn
      else
        let len = xn / wsn in
        let op = Pseudo_operator.Ocopy (ws, Conv.cz_of_int len) in
        Copn(xs,t,Sopn.Opseudo_op op, es)
    | Sopn.Opseudo_op(Ocopy _), _ -> assert false
    | Sopn.Opseudo_op(Pseudo_operator.Oswap _), x::_ ->
      (* Fix the type it is dummy for the moment *)
      let ty = Conv.cty_of_ty (Typing.ty_lval pd loc x) in
      Copn(xs, t, Sopn.Opseudo_op(Pseudo_operator.Oswap ty), es)
    | Sopn.Opseudo_op(Pseudo_operator.Oswap _), [] -> assert false
    | Sopn.Oslh (SLHprotect_ptr _), [Lvar x] ->
      (* Fix the size it is dummy for the moment *)
      let ws, len = array_kind (L.unloc x).v_ty in
      let op = Slh_ops.SLHprotect_ptr (ws, Conv.cz_of_int len) in
      Copn(xs,t, Sopn.Oslh op, es)
    | Sopn.Oslh (SLHprotect_ptr _), _ -> assert false
    | Sopn.Opseudo_op (Odeclassify _), _ ->
      let arg = List.hd es in
      let ty = Conv.cty_of_ty (Typing.ty_expr pd loc arg) in
      Copn(xs, t, Sopn.Opseudo_op (Odeclassify ty), es)
    | Sopn.Opseudo_op (Onop | Odeclassify_mem _ | Omulu _ | Oaddcarry _ | Osubcarry _), _
    | Sopn.Oslh (SLHinit|SLHupdate|SLHmove|SLHprotect _|SLHprotect_ptr_fail _), _
    | Sopn.Oasm _, _ -> ir
    end

  | Csyscall(xs, o, es) ->
    begin match o with
    | Syscall_t.RandomBytes _ ->
      (* Fix the size it is dummy for the moment *)
      let ty =
        match xs with
        | [x] -> Typing.ty_lval pd loc x
        | _ -> assert false in
      let ws, len = array_kind ty in
      Csyscall(xs, Syscall_t.RandomBytes (ws, Conv.cz_of_int len), es)
    end
  | Cassert (msg, e) ->
    Cassert (msg, fix_length_eassert e)

  | Ccall _ -> ir


let fix_length_contract fc =
  let doit = List.map (fun (msg, e) -> (msg, fix_length_eassert e)) in
  {fc with
    f_pre = doit fc.f_pre;
    f_post = doit fc.f_post }

let iac_func pd f =
  { f with f_body = iac_stmt pd f.f_body;
           f_contract = Option.map fix_length_contract f.f_contract}

let doit pd (p:(unit, 'asm) Prog.prog) : (unit, 'asm) Prog.prog =
  (fst p, List.map (iac_func pd) (snd p))