Module RockSaltAsmGen

Require Import Coqlib.
Require Import Integers AST Floats.
Require Import Asm MC.
Require Import X86Model.Encode.
Require Import X86Model.X86Syntax.
Require Import Errors.
Require Import RockSaltAsm.
Require Import Lists.List.
Require Import Segment.
Import ListNotations.

Local Open Scope error_monad_scope.


Definition DMAP_TYPE := ident -> int32.

Definition default_dmap (d:ident) : int32 := Word.zero.

Definition gdef := @FlatAsmProgram.gdef MC.instruction.

Fixpoint transf_globvars (gdefs : list (ident * option gdef * segblock)) (ofs:int32)
  (daccum: list (ident * option (globvar gv_info)))
  : int32 * list (ident * option (globvar gv_info)) :=
  match gdefs with
  | nil => (ofs, daccum)
  | ((id, None, sb) :: gdefs') =>
    transf_globvars gdefs' ofs daccum
  | ((_, Some (Gfun _), sb) :: gdefs') =>
    transf_globvars gdefs' ofs daccum
  | ((id, Some (Gvar v), sb) :: gdefs') =>
    let sz := Word.repr (init_data_list_size (gvar_init v)) in
    let info := mkInfo ofs sz in
    let v' := mkglobvar info v.(gvar_init) v.(gvar_readonly) v.(gvar_volatile) in
    transf_globvars gdefs' (Word.add ofs sz) ((id, Some v')::daccum)
  end.

Definition update_dmap (dmap: DMAP_TYPE) (d:ident) (ofs:int32) : DMAP_TYPE :=
  (fun d' => (if ident_eq d' d then ofs
           else dmap d')).

Fixpoint gvars_to_dmap (l: list (ident * option (globvar gv_info))) : DMAP_TYPE :=
  match l with
  | nil => default_dmap
  | (id, None) :: l' => (gvars_to_dmap l')
  | (id, Some v)::l' =>
    update_dmap (gvars_to_dmap l') id v.(gvar_info).(gv_offset)
  end.




Fixpoint nth_byte (v:int) (n:nat) : int :=
  match n with
  | O =>
    Int.and v (Int.repr 255)
  | S n' =>
    nth_byte (Int.divu v (Int.repr 256)) n'
  end.

Fixpoint nth_byte_64 (v:Integers.int64) (n:nat) : Integers.int64 :=
  match n with
  | O =>
    Int64.and v (Int64.repr 255)
  | S n' =>
    nth_byte_64 (Int64.divu v (Int64.repr 256)) n'
  end.
  
Definition int32_to_bytes (i32:int) :=
  [nth_byte i32 0; nth_byte i32 1; nth_byte i32 2; nth_byte i32 3].

Definition int16_to_bytes (i16:int) :=
  [nth_byte i16 0; nth_byte i16 1].

Definition int64_to_bytes (i64:Integers.int64) :=
  [nth_byte_64 i64 0; nth_byte_64 i64 1; nth_byte_64 i64 2; nth_byte_64 i64 3;
   nth_byte_64 i64 4; nth_byte_64 i64 5; nth_byte_64 i64 6; nth_byte_64 i64 7].

Definition n_zeros (n:nat) : list int8 :=
  List.map (fun _ => Word.zero) (seq 1 n).

Definition transl_init_data (dmap: DMAP_TYPE) (idata: AST.init_data) : ACCUM_DATA :=
  match idata with
  | Init_int8 i => fun data_addr => [Word.repr (Int.unsigned i)]
  | Init_int16 i => fun data_addr => List.map (fun i => Word.repr (Int.unsigned i)) (int16_to_bytes i)
  | Init_int32 i => fun data_addr => List.map (fun i => Word.repr (Int.unsigned i)) (int32_to_bytes i)
  | Init_int64 i => fun data_addr => List.map (fun i => Word.repr (Int64.unsigned i)) (int64_to_bytes i)
  | Init_float32 f => fun data_addr => List.map (fun i => Word.repr (Int64.unsigned i))
                              (int64_to_bytes (Float.to_bits (Float.of_single f)))
  | Init_float64 f => fun data_addr => List.map (fun i => Word.repr (Int64.unsigned i))
                              (int64_to_bytes (Float.to_bits f))
  | Init_space n => fun data_addr => n_zeros (compcert.lib.Coqlib.nat_of_Z n)
  | Init_addrof id ofs =>
    fun data_addr =>
    let i32 := Word.add data_addr (Word.add (dmap id) (Word.repr (Ptrofs.unsigned ofs))) in
    let i32' := Int.repr (Word.unsigned i32) in
    List.map (fun i => Word.repr (Int.unsigned i)) (int32_to_bytes i32')
  end.

Fixpoint collect_init_data (dmap: DMAP_TYPE) (gvars: list (ident * option (globvar gv_info))) : ACCUM_DATA :=
  match gvars with
  | nil => fun data_addr => nil
  | (_,v)::gvars' =>
    let l' := collect_init_data dmap gvars' in
    match v with
    | None => l'
    | Some v =>
      let init_bytes := flat_map_accum_list (transl_init_data dmap) v.(gvar_init) in
      app_accum_list init_bytes l'
    end
  end.


Definition no_prefix : prefix := mkPrefix None None false false.

Definition encode ins := x86_encode no_prefix ins.



Definition transl_ireg (r: ireg) : res register :=
  match r with
  | RAX => OK EAX
  | RBX => OK EBX
  | RCX => OK ECX
  | RDX => OK EDX
  | RSI => OK ESI
  | RDI => OK EDI
  | RBP => OK EBP
  | RSP => OK ESP
  | _ => Error nil
  end.

Definition int_to_bits (i:Int.int) : int32 :=
  (Word.repr (Int.unsigned i)).

Definition int_to_bits8 (i:Int.int) : int8 :=
  (Word.repr (Int.unsigned i)).

Definition ptrofs_to_bits (i:Ptrofs.int) : int32 :=
  (Word.repr (Ptrofs.unsigned i)).


Definition transl_cond_type (t:testcond) :=
  match t with
  | Cond_e => E_ct
  | Cond_ne => NE_ct
  | Cond_b => B_ct
  | Cond_be => BE_ct
  | Cond_ae => NB_ct
  | Cond_a => NBE_ct
  | Cond_l => L_ct
  | Cond_le => LE_ct
  | Cond_ge => NL_ct
  | Cond_g => NLE_ct
  | Cond_p => P_ct
  | Cond_np => NP_ct
  end.


Definition FMAP_TYPE := ident -> int32.
Definition LMAP_TYPE := ident -> label -> int32.

Definition update_lmap (lmap: LMAP_TYPE) (f:ident) (l:label) (ofs:int32) : LMAP_TYPE :=
  (fun f' l' =>
     if ident_eq f' f then
       (if ident_eq l' l then ofs
        else lmap f' l')
     else lmap f' l').

Definition update_fmap (fmap: FMAP_TYPE) (f:ident) (ofs:int32) : FMAP_TYPE :=
  (fun f' =>
     if ident_eq f' f then ofs
     else fmap f').


Definition instr_size (i:instr) : res int32 :=
  match (encode i) with
  | None => Error (msg (String.append (instr_to_string i) " instruction has no encoding"))
  | Some l => OK (Word.repr (Z.of_nat (length l)))
  end.

Fixpoint instr_list_size (il:list instr) : res int32 :=
  match il with
  | nil => OK Word.zero
  | i::il' =>
    do isz <- instr_size i;
    do ilsz <- instr_list_size il';
    OK (Word.add isz ilsz)
  end.


Section TRANSL.
Variable dmap : DMAP_TYPE.


Definition transl_scale (s:Z) : res scale :=
  match s with
  | 1 => OK Scale1
  | 2 => OK Scale2
  | 4 => OK Scale4
  | 8 => OK Scale8
  | _ => Error (msg "Translation of scale failed")
  end.

Definition transl_addr_mode (m: addrmode) : res (int32 -> address) :=
  match m with
  | Addrmode b ofs const =>
    do base_reg <-
         match b with
         | None => OK None
         | Some r =>
           do r' <- transl_ireg r; OK (Some r')
         end;
    do index <-
         match ofs with
         | None => OK None
         | Some (r,scale) =>
           do r' <- transl_ireg r;
           do sc <- transl_scale scale;
             OK (Some (sc, r'))
         end;
    OK (fun data_addr =>
      let disp :=
        match const with
        | inl disp' => (Word.repr disp')
        | inr (ident,ptrofs) =>
          (Word.add data_addr (Word.add (dmap ident) (ptrofs_to_bits ptrofs)))
        end in
      {|
        addrDisp := disp;
        addrBase := base_reg;
        addrIndex := index;
      |})
  end.





Definition transl_instr (fmap: FMAP_TYPE) (lmap: LMAP_TYPE)
           (f:ident) (i: MC.instruction) (ofs: int32) :
  res (ACCUM_INSTRS * LMAP_TYPE) :=
  match i with
  | MCjmp_l l ofs' =>
    do isz <- instr_size (JMP true false (Imm_op Word.zero) None);
    let rel_ofs := Word.sub (lmap f l) (Word.add ofs isz) in
    OK (fun data_addr => [JMP true false (Imm_op rel_ofs) None],lmap)
  | MCjcc c l ofs' =>
    do isz <- instr_size (Jcc (transl_cond_type c) Word.zero);
    let rel_ofs := Word.sub (lmap f l) (Word.add ofs isz) in
    OK (fun data_addr => [Jcc (transl_cond_type c) rel_ofs],lmap)
  | MCAsminstr i =>
    match i with
    | Pxorl_r rd =>
      do rd' <- transl_ireg rd;
        OK (fun data_addr => [XOR true (Reg_op rd') (Reg_op rd')],lmap)
    | Paddl_ri rd n =>
      do rd' <- transl_ireg rd;
        OK (fun data_addr => [ADD true (Reg_op rd') (Imm_op (int_to_bits n))], lmap)
    | Psubl_ri rd n =>
      do rd' <- transl_ireg rd;
        OK (fun data_addr => [SUB true (Reg_op rd') (Imm_op (int_to_bits n))], lmap)
    | Psubl_rr rd r1 =>
      do rd' <- transl_ireg rd;
        do r1' <- transl_ireg r1;
        OK (fun data_addr => [SUB true (Reg_op rd') (Reg_op r1')], lmap)
    | Pleal rd a =>
      do rd' <- transl_ireg rd;
        do a' <- transl_addr_mode a;
        OK (fun data_addr => [LEA (Reg_op rd') (Address_op (a' data_addr))], lmap)
    | Pmovl_ri rd n =>
      do rd' <- transl_ireg rd;
        OK (fun data_addr => [MOV true (Reg_op rd') (Imm_op (int_to_bits n))], lmap)
    | Pmov_rr rd r1 =>
      do rd' <- transl_ireg rd;
        do r1' <- transl_ireg r1;
        OK (fun data_addr => [MOV true (Reg_op rd') (Reg_op r1')],lmap)
    | Pmovl_mr a rs =>
      do a' <- transl_addr_mode a;
        do rs' <- transl_ireg rs;
        OK (fun data_addr => [MOV true (Address_op (a' data_addr)) (Reg_op rs')],lmap)
    | Pmov_mr_a a rs =>
      do a' <- transl_addr_mode a;
        do rs' <- transl_ireg rs;
        OK (fun data_addr => [MOV true (Address_op (a' data_addr)) (Reg_op rs')],lmap)
    | Pmovl_rm rd a =>
      do rd' <- transl_ireg rd;
        do a' <- transl_addr_mode a;
        OK (fun data_addr => [MOV true (Reg_op rd') (Address_op (a' data_addr))],lmap)
    | Pmov_rm_a rd a =>
      do rd' <- transl_ireg rd;
        do a' <- transl_addr_mode a;
        OK (fun data_addr => [MOV true (Reg_op rd') (Address_op (a' data_addr))],lmap)
    | Ptestl_rr r1 r2 =>
      do r1' <- transl_ireg r1;
        do r2' <- transl_ireg r2;
        OK (fun data_addr => [TEST true (Reg_op r1') (Reg_op r2')],lmap)
    | Pcall ros _ =>
      match ros with
      | inr symb =>
        do isz <- instr_size (CALL true false (Imm_op Word.zero) None);
          let rel_ofs := Word.sub (fmap symb) (Word.add ofs isz) in
          OK (fun data_addr => [CALL true false (Imm_op rel_ofs) None],lmap)
      | inl r =>
        Error (MSG (Asm.instr_to_string i) :: MSG " not supported" :: nil)
      end
    | Pret =>
      OK (fun data_addr => [RET true None],lmap)
    | Pimull_rr rd r1 =>
      do rd' <- transl_ireg rd;
        do r1' <- transl_ireg r1;
        OK (fun data_addr => [IMUL false (Reg_op rd') (Some (Reg_op r1')) None],lmap)
    | Pimull_ri rd n =>
      do rd' <- transl_ireg rd;
        OK (fun data_addr => [IMUL true (Reg_op rd') (Some (Reg_op rd')) (Some (int_to_bits n))],lmap)
    | Pcmpl_rr r1 r2 =>
      do r1' <- transl_ireg r1;
        do r2' <- transl_ireg r2;
        OK (fun data_addr => [CMP true (Reg_op r1') (Reg_op r2')],lmap)
    | Pcmpl_ri r1 n =>
      do r1' <- transl_ireg r1;
        OK (fun data_addr => [CMP true (Reg_op r1') (Imm_op (int_to_bits n))],lmap)
    | Pcltd =>
      OK (fun data_addr => [CDQ], lmap)
    | Pidivl r1 =>
      do r1' <- transl_ireg r1;
        OK (fun data_addr => [IDIV true (Reg_op r1')],lmap)
    | Psall_ri rd n =>
      do rd' <- transl_ireg rd;
        OK (fun data_addr => [SHL true (Reg_op rd') (Imm_ri (int_to_bits8 n))], lmap)
    | Plabel l =>
      OK (fun data_addr => [FNOP], update_lmap lmap f l ofs)
    | _ =>
      Error (MSG (Asm.instr_to_string i) :: MSG " not supported" :: nil)
    end
  | _ =>
    Error (MSG (MC.instr_to_string i) :: MSG " not supported" :: nil)
  end.

Fixpoint transl_instr_list (fmap: FMAP_TYPE) (lmap: LMAP_TYPE)
           (f:ident) (il: list MC.instruction) (ofs: int32) (accum : ACCUM_INSTRS)
  : res (ACCUM_INSTRS * LMAP_TYPE * int32) :=
  match il with
  | nil => OK (accum, lmap, ofs)
  | i::il' =>
    do (instrs, lmap') <-
       transl_instr fmap lmap f i ofs;
    do isz <- instr_list_size (instrs Word.zero);
    transl_instr_list fmap lmap' f il' (Word.add ofs isz)
                      (app_accum_list instrs accum)
  end.

Definition transl_function (fmap: FMAP_TYPE) (lmap: LMAP_TYPE)
         (f:ident) (code:list MC.instruction) (ofs:int32) (accum: ACCUM_INSTRS)
  : res (ACCUM_INSTRS * FMAP_TYPE * LMAP_TYPE * int32) :=
  do (im, ofs') <- transl_instr_list fmap lmap f code ofs accum;
  let (accum', lmap') := im in
  OK (accum', update_fmap fmap f ofs, lmap', ofs').

Definition transl_function' (fmap: FMAP_TYPE) (lmap: LMAP_TYPE)
           (f:ident) (code: list MC.instr_with_info) (ofs:int32) (accum: ACCUM_INSTRS) :=
  let code' := List.map (fun '(i,_,_) => i) code in
  transl_function fmap lmap f code' ofs accum.


Fixpoint transf_globfuns (fmap: FMAP_TYPE) (lmap: LMAP_TYPE)
         (gdefs : list (ident * option gdef * segblock)) (ofs:int32)
  (iaccum: ACCUM_INSTRS) (daccum: list (ident * option RockSaltAsm.fundef))
  : res (ACCUM_INSTRS * list (ident * option RockSaltAsm.fundef) * FMAP_TYPE * LMAP_TYPE * int32) :=
  match gdefs with
  | nil => OK (iaccum, daccum, fmap, lmap, ofs)
  | ((id, None, sb) :: gdefs') =>
    transf_globfuns fmap lmap gdefs' ofs iaccum daccum
  | ((id, Some (Gvar _), sb) :: gdefs') =>
    transf_globfuns fmap lmap gdefs' ofs iaccum daccum
  | ((id, Some (Gfun fn), sb) :: gdefs') =>
    match fn with
    | Internal f =>
      do r <- transl_function' fmap lmap id (@FlatAsmProgram.fn_code MC.instruction f) ofs iaccum;
      let '(il, fmap', lmap', ofs') := r in
      let f' := mkFun ofs (Word.sub ofs' ofs) in
      transf_globfuns fmap' lmap' gdefs' ofs'
                      il
                      ((id, Some (Internal f')) :: daccum)
    | External ef =>
      transf_globfuns fmap lmap gdefs' ofs iaccum
                      ((id, Some (External ef)) :: daccum)
    end
  end.

End TRANSL.

  

Definition default_fmap (f:ident) : int32 := Word.zero.
Definition default_lmap (f:ident) (l:label) : int32 := Word.zero.

Definition align_int32 (n: int32) (amount: Z) :int32 :=
  Word.repr (align (Word.unsigned n) amount).

Definition prog_defns := @FlatAsmProgram.prog_defs MC.instruction.

Definition transf_program (p: MC.program) : res RockSaltAsm.program :=
  let (data_seg_size, gvars) :=
      transf_globvars (prog_defns p) Word.zero [] in
  let dmap := gvars_to_dmap gvars in
  let init_data := collect_init_data dmap gvars in
  do r <- transf_globfuns dmap
             default_fmap default_lmap
             (prog_defns p) Word.zero (fun data_addr => []) [];
    let '(_,_,fmap,lmap,_) := r in
  do r <- transf_globfuns dmap
             fmap lmap
             (prog_defns p) Word.zero (fun data_addr => []) [];
    let '(_,_,fmap,lmap,_) := r in
  do r <- transf_globfuns dmap
             fmap lmap
             (prog_defns p) Word.zero (fun data_addr => []) [];
  let '(instrs,gfuns,fmap,_,code_seg_size) := r in
  let gvars_defs :=
      List.map (fun (e: ident * option (globvar gv_info)) =>
                  let (id,gv) := e in
                  (id, match gv with
                       | None => None
                       | Some gv => Some (Gvar gv)
                       end))
               gvars in
  let gfuns_defs :=
      List.map (fun (e: ident * option fundef) =>
                  let (id,fn) := e in
                  (id, match fn with
                       | None => None
                       | Some fn => Some (Gfun fn)
                       end))
               gfuns in
  let tinstrs := fun data_addr => List.rev (instrs data_addr) in
  OK {|
    prog_defs := gvars_defs ++ gfuns_defs;
    prog_public := @FlatAsmProgram.prog_public MC.instruction p;
    prog_main := @FlatAsmProgram.prog_main MC.instruction p;
    text_instrs := tinstrs;
    init_data := init_data;
    main_ofs := fmap (@FlatAsmProgram.prog_main MC.instruction p);
  |}.