Require Import Coqlib Integers AST Maps.
Require Import Asm FlatAsm Segment.
Require Import Errors.
Require Import FlatAsmBuiltin.
Require Import Memtype.
Require Import FlatAsmProgram.
Import ListNotations.


Local Open Scope error_monad_scope.


Definition alignw:Z := 8.


Definition data_label (ofs:Z) : seglabel := (data_segid, Ptrofs.repr ofs).
Definition code_label (ofs:Z) : seglabel := (code_segid, Ptrofs.repr ofs).

Definition default_gid_map : GID_MAP_TYPE := fun id => None.
Definition default_label_map : LABEL_MAP_TYPE :=
  fun id l => None.

Definition update_gid_map (id:ident) (l:seglabel) (map:GID_MAP_TYPE) : GID_MAP_TYPE :=
  fun id' => if peq id id' then Some l else (map id').

Definition update_label_map (id:ident) (l:Asm.label) (tl:seglabel) (map:LABEL_MAP_TYPE) : LABEL_MAP_TYPE :=
  fun id' l' => if peq id id' then (if peq l l' then Some tl else (map id' l')) else (map id' l').


Section WITH_GID_MAP.

Variable gid_map : GID_MAP_TYPE.













Section WITH_LABEL_MAP.
Variable label_map : LABEL_MAP_TYPE.





Definition transl_instr (ofs:Z) (fid: ident) (sid:segid_type) (i:Asm.instruction) : res FlatAsm.instr_with_info :=
  match i with
      Pallocframe _ _ _
    | Pfreeframe _ _
    | Pload_parent_pointer _ _ => Error (MSG "No pseudo instructions" :: nil)
    | _ =>
      let sz := instr_size i in
      let sblk := mkSegBlock sid (Ptrofs.repr ofs) (Ptrofs.repr sz) in
      OK (i, sblk, fid)
  end.

Fixpoint transl_instrs (fid:ident) (sid:segid_type) (ofs:Z) (instrs: list Asm.instruction) : res (Z * list instr_with_info) :=
  match instrs with
  | nil => OK (ofs, nil)
  | i::instrs' =>
    let sz := instr_size i in
    let nofs := ofs+sz in
    do instr <- transl_instr ofs fid sid i;
    do (fofs, tinstrs') <- transl_instrs fid sid nofs instrs';
    OK (fofs, instr :: tinstrs')
  end.

Definition transl_fun (fid: ident) (f:Asm.function) : res function :=
  match gid_map fid with
  | None => Error (MSG "Translation of function fails: no address for this function" :: nil)
  | Some (sid, ofs) =>
    let ofs' := Ptrofs.unsigned ofs in
    do (fofs, code') <- transl_instrs fid sid ofs' (Asm.fn_code f);
      if zle fofs Ptrofs.max_unsigned then
        (let sz := (Asm.code_size (Asm.fn_code f)) in
         let sblk := mkSegBlock sid ofs (Ptrofs.repr sz) in
         OK (mkfunction (Asm.fn_sig f) code' sblk (Asm.fn_stacksize f) (Asm.fn_pubrange f)))
      else
        Error (MSG "The size of the function exceeds limit" ::nil)
  end.


Definition transl_globdef (id:ident) (def: option (AST.globdef Asm.fundef unit))
  : res (ident * option FlatAsmProgram.gdef * segblock) :=
  match def with
  | None => OK (id, None, (mkSegBlock undef_segid Ptrofs.zero Ptrofs.zero))
  | Some (AST.Gvar v) =>
    match gid_map id with
    | None => Error (MSG "Translation of a global variable fails: no address for the variable" :: nil)
    | Some (sid,ofs) =>
      let sz := AST.init_data_list_size (AST.gvar_init v) in
      let sblk := mkSegBlock sid ofs (Ptrofs.repr sz) in
      OK (id, Some (Gvar v), sblk)
    end
  | Some (AST.Gfun f) =>
    match f with
    | External f =>
      let sblk := mkSegBlock undef_segid Ptrofs.zero Ptrofs.zero in
      OK (id, Some (Gfun (External f)), sblk)
    | Internal fd =>
      do fd' <- transl_fun id fd;
        OK (id, Some (Gfun (Internal fd')), (fn_range fd'))
    end
  end.

Fixpoint transl_globdefs (defs : list (ident * option (AST.globdef Asm.fundef unit)))
  : res (list (ident * option gdef * segblock)) :=
  match defs with
  | nil => OK nil
  | (id, def)::defs' =>
    do tdef <- transl_globdef id def;
    do tdefs' <- transl_globdefs defs';
    OK (tdef :: tdefs')
  end.
  
Definition transl_prog_with_map (p:Asm.program) (data_sz code_sz:Z): res program :=
  do defs <- transl_globdefs (AST.prog_defs p);
  OK (Build_program
        defs
        (AST.prog_public p)
        (AST.prog_main p)
        (mkSegment data_segid (Ptrofs.repr data_sz))
        (mkSegment code_segid (Ptrofs.repr code_sz))
        gid_map
        label_map
        (Globalenvs.Genv.to_senv (Globalenvs.Genv.globalenv p)))
      .


End WITH_LABEL_MAP.

End WITH_GID_MAP.

Compute mappings from global identifiers to section labels

Definition is_label (i: Asm.instruction) : option Asm.label :=
  match i with
  | Asm.Plabel l => Some l
  | _ => None
  end.

Definition update_instr (lmap: ident -> Asm.label -> option seglabel) (csize: Z) (fid: ident) (i: Asm.instruction) :=
  let csize' := csize + instr_size i in
  let lmap' :=
      match is_label i with
      | Some l => update_label_map fid l (code_label csize') lmap
      | _ => lmap
      end in
  (lmap', csize').

Definition update_instrs lmap csize fid c : LABEL_MAP_TYPE * Z :=
  fold_left (fun (ls : LABEL_MAP_TYPE * Z) i =>
               let '(lmap, csize) := ls in
               update_instr lmap csize fid i) c (lmap, csize).

Definition update_maps_def
           (gmap: ident -> option seglabel)
           (lmap: ident -> Asm.label -> option seglabel)
           (dsize csize: Z) (i: ident) (def: option (AST.globdef Asm.fundef unit)):
  (GID_MAP_TYPE * LABEL_MAP_TYPE * Z * Z) :=
  match def with
  | None => (gmap, lmap, dsize, csize)
  | Some (AST.Gvar gvar) =>
    let sz:= AST.init_data_list_size (AST.gvar_init gvar) in
    (update_gid_map i (data_label dsize) gmap, lmap, dsize + align sz alignw, csize)
  | Some (AST.Gfun (External ef)) => (gmap, lmap, dsize, csize)
  | Some (AST.Gfun (Internal f)) =>
    let (lmap', csize') := update_instrs lmap csize i (Asm.fn_code f) in
    (update_gid_map i (code_label csize) gmap, lmap', dsize, align csize' alignw)
  end.

Definition update_maps gmap lmap dsize csize defs : (GID_MAP_TYPE * LABEL_MAP_TYPE * Z * Z) :=
  fold_left (fun (acc : GID_MAP_TYPE * LABEL_MAP_TYPE * Z * Z)
               (id: ident * option (AST.globdef Asm.fundef unit)) =>
               let '(gmap, lmap, dsize, csize) := acc in
               let '(i,def) := id in
               update_maps_def gmap lmap dsize csize i def)
            defs (gmap, lmap, dsize, csize).

Definition make_maps (p:Asm.program) : (GID_MAP_TYPE * LABEL_MAP_TYPE * Z * Z) :=
  update_maps default_gid_map default_label_map 0 0 (AST.prog_defs p).



Definition no_duplicated_defs {F V: Type} (defs: list (ident * option (AST.globdef F V))) :=
  list_norepet (map fst defs).

Fixpoint labels (c: Asm.code) : list Asm.label :=
  match c with
  | [] => []
  | i :: r =>
    match is_label i with
      Some l => l :: labels r
    | None => labels r
    end
  end.

Definition no_duplicated_labels (c: Asm.code) :=
  list_norepet (labels c).

Definition globdef_no_duplicated_labels (def: option (AST.globdef Asm.fundef unit)) :=
  match def with
  | Some (AST.Gfun (Internal f)) => no_duplicated_labels (Asm.fn_code f)
  | _ => True
  end.

Definition globdef_no_duplicated_labels_dec def : { globdef_no_duplicated_labels def } + { ~ globdef_no_duplicated_labels def }.

Definition defs_no_duplicated_labels (defs: list (ident * _)) :=
  Forall globdef_no_duplicated_labels (map snd defs).

Definition def_size (def: AST.globdef Asm.fundef unit) : Z :=
  match def with
  | AST.Gfun (External e) => 1
  | AST.Gfun (Internal f) => Asm.code_size (Asm.fn_code f)
  | AST.Gvar v => AST.init_data_list_size (AST.gvar_init v)
  end.

Definition odef_size (def: option (AST.globdef Asm.fundef unit)) : Z :=
  match def with
  | Some def => def_size def
  | _ => 0
  end.

Lemma def_size_pos:
  forall d,
    0 <= def_size d.

Lemma odef_size_pos:
  forall d,
    0 <= odef_size d.

Definition def_not_empty def : Prop :=
  match def with
  | None => True
  | Some def' => 0 < def_size def'
  end.


Definition defs_not_empty defs :=
  Forall def_not_empty defs.

Definition defs_not_empty_dec defs : { defs_not_empty defs } + { ~ defs_not_empty defs }.

Definition main_exists main (defs: list (ident * option (AST.globdef Asm.fundef unit))) :=
  Exists (fun '(id, def) =>
            main = id
            /\ match def with
              | None => False
              | Some _ => True
              end) defs.

Definition main_exists_dec main defs : {main_exists main defs } + {~ main_exists main defs}.

Definition has_label (i: Asm.instruction) (lab: label) :=
  match i with
    Pjmp_l l
  | Pjcc _ l
  | Pjcc2 _ _ l => l = lab
  | Pjmptbl _ l => In lab l
  | _ => False
  end.

Definition valid_labels_funs (f: Asm.function) :=
  Forall
    (fun i : Asm.instruction =>
       forall lab : label, has_label i lab -> In lab (labels (Asm.fn_code f)))
    (Asm.fn_code f).

Definition valid_labels_funs_dec f: {valid_labels_funs f} + { ~ valid_labels_funs f}.

Definition valid_labels (defs: list (ident * option (globdef Asm.fundef unit))) :=
  Forall (fun '(i,d) =>
            match d with
              Some (Gfun (Internal f)) => valid_labels_funs f
            | _ => True
            end
         ) defs.

Definition valid_labels_dec defs: {valid_labels defs} + { ~ valid_labels defs}.

Record wf_prog (p:Asm.program) : Prop :=
  {
    wf_prog_not_empty: defs_not_empty (map snd (AST.prog_defs p));
    wf_prog_norepet_defs: list_norepet (map fst (AST.prog_defs p));
    wf_prog_norepet_labels: defs_no_duplicated_labels (AST.prog_defs p);
    wf_prog_main_exists: main_exists (AST.prog_main p) (AST.prog_defs p);
    wf_prog_valid_labels: valid_labels (AST.prog_defs p);
  }.

Definition check_wellformedness p : { wf_prog p } + { ~ wf_prog p }.

Definition transf_program (p:Asm.program) : res program :=
  if check_wellformedness p then
    (let r := make_maps p in
     let '(gmap,lmap,dsize,csize) := r in
     if zle (dsize + csize) Ptrofs.max_unsigned then
       transl_prog_with_map gmap lmap p dsize csize
     else
       Error (MSG "Size of segments too big" :: nil))
  else
    Error (MSG "Program not well-formed. There exists duplicated labels or definitions" :: nil).