include "utilities/lists.ma".
include "common/Globalenvs.ma".
include "Cminor/syntax.ma".
include "RTLabs/syntax.ma".

definition env ≝ identifier_map SymbolTag register.
definition label_env ≝ identifier_map SymbolTag label.
definition populate_env : env → universe RegisterTag → list (ident × typ) → list (register×typ) × env × (universe RegisterTag) ≝
λen,gen. foldr ??
 (λidt,rsengen.
   let 〈id,ty〉 ≝ idt in
   let 〈rs,en,gen〉 ≝ rsengen in
   let 〈r,gen'〉 ≝ fresh … gen in
     〈〈r,ty〉::rs, add ?? en id r, gen'〉) 〈[ ], en, gen〉.

definition populate_label_env : label_env → universe LabelTag → list ident → label_env × (universe LabelTag) ≝
λen,gen. foldr ??
 (λid,engen.
   let 〈en,gen〉 ≝ engen in
   let 〈r,gen'〉 ≝ fresh … gen in
     〈add ?? en id r, gen'〉) 〈en, gen〉.

lemma lookup_sigma : ∀tag,A,m. ∀i:(Σl:identifier tag. lookup tag A m l ≠ None ?).
  lookup tag A m i ≠ None ?.
#tag #A #m * #i #H @H
qed.

(* Add a statement to the graph, *without* updating the entry label. *)
definition fill_in_statement : label → statement → internal_function → internal_function ≝
λl,s,f.
  mk_internal_function (f_labgen f) (f_reggen f)
                       (f_result f) (f_params f) (f_locals f)
                       (f_stacksize f) (add ?? (f_graph f) l s)
                       (dp ?? (f_entry f) ?) (dp ?? (f_exit f) ?).
@lookup_add_oblivious @lookup_sigma
qed.

(* Add a statement to the graph, making it the entry label. *)
definition add_to_graph : label → statement → internal_function → internal_function ≝
λl,s,f.
  mk_internal_function (f_labgen f) (f_reggen f)
                       (f_result f) (f_params f) (f_locals f)
                       (f_stacksize f) (add ?? (f_graph f) l s)
                       (dp ?? l ?) (dp ?? (f_exit f) ?).
[ >lookup_add_hit % #E destruct
| @lookup_add_oblivious @lookup_sigma
] qed.

(* Add a statement with a fresh label to the start of the function.  The
   statement is parametrised by the *next* instruction's label. *)
definition add_fresh_to_graph : (label → statement) → internal_function → internal_function ≝
λs,f.
  let 〈l,g〉 ≝ fresh … (f_labgen f) in
  let s' ≝ s (f_entry f) in
    (mk_internal_function g (f_reggen f)
                       (f_result f) (f_params f) (f_locals f)
                       (f_stacksize f) (add ?? (f_graph f) l s')
                       (dp ?? l ?) (dp ?? (f_exit f) ?)).
[ >lookup_add_hit % #E destruct
| @lookup_add_oblivious @lookup_sigma
] qed.

definition fresh_reg : internal_function → typ → register × internal_function ≝
λf,ty.
  let 〈r,g〉 ≝ fresh … (f_reggen f) in
    〈r, mk_internal_function (f_labgen f) g
                       (f_result f) (f_params f) (〈r,ty〉::(f_locals f))
                       (f_stacksize f) (f_graph f) (f_entry f) (f_exit f)〉.

axiom UnknownVariable : String.

definition choose_reg : env → ∀ty.expr ty → internal_function → res (register × internal_function) ≝
λenv,ty,e,f.
  match e with
  [ Id _ i ⇒
      do r ← opt_to_res … [MSG UnknownVariable; CTX ? i] (lookup … env i);
      OK ? 〈r, f〉
  | _ ⇒
      OK ? (fresh_reg f ty)
  ].

definition choose_regs : env → list (Σt:typ.expr t) → internal_function → res (list register × internal_function) ≝
λenv,es,f.
  foldr (Σt:typ.expr t) ?
    (λe,acc. do 〈rs,f〉 ← acc;
             do 〈r,f'〉 ← match e with [ dp t e ⇒ choose_reg env t e f ];
             OK ? 〈r::rs,f'〉) (OK ? 〈[ ], f〉) es.

axiom BadCminorProgram : String.

let rec add_expr (env:env) (ty:typ) (e:expr ty) (dst:register) (f:internal_function) on e: res internal_function ≝
match e with
[ Id _ i ⇒
    do r ← opt_to_res … [MSG UnknownVariable; CTX ? i] (lookup ?? env i);
    match register_eq r dst with
    [ inl _ ⇒ OK ? f
    | inr _ ⇒ OK ? (add_fresh_to_graph (St_op1 Oid dst r) f)
    ]
| Cst _ c ⇒ OK ? (add_fresh_to_graph (St_const dst c) f)
| Op1 _ _ op e' ⇒
    do 〈r,f〉 ← choose_reg env ? e' f;
    let f ≝ add_fresh_to_graph (St_op1 op dst r) f in
    add_expr env ? e' r f
| Op2 _ _ _ op e1 e2 ⇒
    do 〈r1,f〉 ← choose_reg env ? e1 f;
    do 〈r2,f〉 ← choose_reg env ? e2 f;
    let f ≝ add_fresh_to_graph (St_op2 op dst r1 r2) f in
    do f ← add_expr env ? e2 r2 f;
    add_expr env ? e1 r1 f
| Mem _ _ q e' ⇒
    do 〈r,f〉 ← choose_reg env ? e' f;
    let f ≝ add_fresh_to_graph (St_load q r dst) f in
    add_expr env ? e' r f
| Cond _ _ _ e' e1 e2 ⇒
    let resume_at ≝ f_entry f in
    do f ← add_expr env ? e2 dst f;
    let lfalse ≝ f_entry f in
    let f ≝ add_fresh_to_graph (λ_.St_skip resume_at) f in
    do f ← add_expr env ? e1 dst f;
    do 〈r,f〉 ← choose_reg env ? e' f;
    let f ≝ add_fresh_to_graph (λltrue. St_cond r ltrue lfalse) f in
    add_expr env ? e' r f
| Ecost _ l e' ⇒
    do f ← add_expr env ? e' dst f;
    OK ? (add_fresh_to_graph (St_cost l) f)
].

(* This shouldn't occur; maybe use vectors? *)
axiom WrongNumberOfArguments : String.

let rec add_exprs (env:env) (es:list (Σt:typ.expr t)) (dsts:list register) (f:internal_function) on es: res internal_function ≝
match es with
[ nil ⇒ match dsts with [ nil ⇒ OK ? f | cons _ _ ⇒ Error ? (msg WrongNumberOfArguments) ]
| cons e et ⇒
    match dsts with
    [ nil ⇒ Error ? (msg WrongNumberOfArguments)
    | cons r rt ⇒
        do f ← add_exprs env et rt f;
        match e with [ dp ty e ⇒ add_expr env ty e r f ]
    ]
].

axiom UnknownLabel : String.
axiom ReturnMismatch : String.

let rec add_stmt (env:env) (label_env:label_env) (s:stmt) (exits:list label) (f:internal_function) on s : res internal_function ≝
match s with
[ St_skip ⇒ OK ? f
| St_assign _ x e ⇒
    do dst ← opt_to_res … [MSG UnknownVariable; CTX ? x] (lookup ?? env x);
    add_expr env ? e dst f
| St_store _ _ q e1 e2 ⇒
    do 〈val_reg, f〉 ← choose_reg env ? e2 f;
    do 〈addr_reg, f〉 ← choose_reg env ? e1 f;
    let f ≝ add_fresh_to_graph (St_store q addr_reg val_reg) f in
    do f ← add_expr env ? e1 addr_reg f;
    add_expr env ? e2 val_reg f
| St_call return_opt_id e args ⇒
    do return_opt_reg ←
      match return_opt_id with
      [ None ⇒ OK ? (None ?)
      | Some id ⇒ do r ← opt_to_res … [MSG UnknownVariable; CTX ? id] (lookup ?? env id); OK ? (Some ? r)
      ];
    do 〈args_regs, f〉 ← choose_regs env args f;
    do f ←
      match e with
      [ Id _ id ⇒ OK ? (add_fresh_to_graph (St_call_id id args_regs return_opt_reg) f)
      | _ ⇒ 
        do 〈fnr, f〉 ← choose_reg env ? e f;
        let f ≝ add_fresh_to_graph (St_call_ptr fnr args_regs return_opt_reg) f in
        add_expr env ? e fnr f
      ];
    add_exprs env args args_regs f
| St_tailcall e args ⇒
    do 〈args_regs, f〉 ← choose_regs env args f;
    do f ←
      match e with
      [ Id _ id ⇒ OK ? (add_fresh_to_graph (λ_. St_tailcall_id id args_regs) f)
      | _ ⇒ 
        do 〈fnr, f〉 ← choose_reg env ? e f;
        let f ≝ add_fresh_to_graph (λ_. St_tailcall_ptr fnr args_regs) f in
        add_expr env ? e fnr f
      ];
    add_exprs env args args_regs f
| St_seq s1 s2 ⇒
    do f ← add_stmt env label_env s2 exits f;
    add_stmt env label_env s1 exits f
| St_ifthenelse _ _ e s1 s2 ⇒ 
    let l_next ≝ f_entry f in
    do f ← add_stmt env label_env s2 exits f;
    let l2 ≝ f_entry f in
    let f ≝ add_fresh_to_graph ? (* XXX: fails, but works if applied afterwards: λ_. St_skip l_next*) f in
    do f ← add_stmt env label_env s1 exits f;
    do 〈r,f〉 ← choose_reg env ? e f;
    let f ≝ add_fresh_to_graph (λl1. St_cond r l1 l2) f in
    add_expr env ? e r f
| St_loop s ⇒
    let f ≝ add_fresh_to_graph ? f in (* dummy statement, fill in afterwards *)
    let l_loop ≝ f_entry f in
    do f ← add_stmt env label_env s exits f;
    OK ? (fill_in_statement l_loop (* XXX another odd failure: St_skip (f_entry f)*)? f)
| St_block s ⇒
    add_stmt env label_env s ((f_entry f)::exits) f
| St_exit n ⇒
    do l ← opt_to_res … (msg BadCminorProgram) (nth_opt ? n exits);
    OK ? (add_fresh_to_graph (* XXX another: λ_. St_skip l*)? f)
| St_switch sz sg e tab n ⇒
    do 〈r,f〉 ← choose_reg env ? e f;
    do l_default ← opt_to_res … (msg BadCminorProgram) (nth_opt ? n exits);
    let f ≝ add_fresh_to_graph (* XXX grrrr: λ_. St_skip l_default*)? f in
    do f ← foldr ?? (λcs,f.
      do f ← f;
      let 〈i,n〉 ≝ cs in
      let 〈cr,f〉 ≝ fresh_reg … f (ASTint sz sg) in
      let 〈br,f〉 ≝ fresh_reg … f (ASTint I8 Unsigned) in
      do l_case ← opt_to_res … (msg BadCminorProgram) (nth_opt ? n exits);
      let f ≝ add_fresh_to_graph (St_cond br l_case) f in
      let f ≝ add_fresh_to_graph (St_op2 (Ocmpu Ceq) (* signed? *) br r cr) f in
        OK ? (add_fresh_to_graph (St_const cr (Ointconst ? i)) f)) (OK ? f) tab;
    add_expr env ? e r f
| St_return opt_e ⇒
    let f ≝ add_fresh_to_graph (λ_. St_return) f in
    match opt_e with
    [ None ⇒ OK ? f
    | Some e ⇒
        match f_result f with
        [ None ⇒ Error ? (msg ReturnMismatch)
        | Some r ⇒ match e with [ dp ty e ⇒ add_expr env ? e (\fst r) f ]
        ]
    ]
| St_label l s' ⇒
    do f ← add_stmt env label_env s' exits f;
    do l' ← opt_to_res … [MSG UnknownLabel; CTX ? l] (lookup ?? label_env l);
    OK ? (add_to_graph l' (* XXX again: St_skip (f_entry f)*)? f)
| St_goto l ⇒
    do l' ← opt_to_res … [MSG UnknownLabel; CTX ? l] (lookup ?? label_env l);
    OK ? (add_fresh_to_graph (* XXX again: λ_.St_skip l'*)? f)
| St_cost l s' ⇒
    do f ← add_stmt env label_env s' exits f;
    OK ? (add_fresh_to_graph (St_cost l) f)
].
[ @(λ_. St_skip l_next)
| @(St_skip (f_entry f))
| @St_skip
| @(λ_. St_skip l)
| @(λ_. St_skip l_default)
| @(St_skip (f_entry f))
| @(λ_.St_skip l')
] qed. 

(* Get labels from a Cminor statement. *)
let rec labels_of (s:stmt) : list ident ≝
match s with
[ St_seq s1 s2 ⇒ (labels_of s1) @ (labels_of s2)
| St_ifthenelse _ _ _ s1 s2 ⇒ (labels_of s1) @ (labels_of s2)
| St_loop s ⇒ labels_of s
| St_block s ⇒ labels_of s
| St_label l s ⇒ l::(labels_of s)
| St_cost _ s ⇒ labels_of s
| _ ⇒ [ ]
].

definition c2ra_function (*: internal_function → internal_function*) ≝
λf.
let labgen0 ≝ new_universe LabelTag in
let reggen0 ≝ new_universe RegisterTag in
let cminor_labels ≝ labels_of (f_body f) in
let 〈params, env1, reggen1〉 ≝ populate_env (empty_map …) reggen0 (f_params f) in
let 〈locals0, env, reggen2〉 ≝ populate_env env1 reggen1 (f_vars f) in
let 〈result, locals, reggen〉 ≝
  match f_return f with
  [ None ⇒ 〈None ?, locals0, reggen2〉
  | Some ty ⇒
      let 〈r,gen〉 ≝ fresh … reggen2 in
        〈Some ? 〈r,ty〉, 〈r,ty〉::locals0, gen〉 ] in
let 〈label_env, labgen1〉 ≝ populate_label_env (empty_map …) labgen0 cminor_labels in
let 〈l,labgen〉 ≝ fresh … labgen1 in
let emptyfn ≝
  mk_internal_function
    labgen
    reggen
    result
    params
    locals
    (f_stacksize f)
    (add ?? (empty_map …) l St_return)
    l
    l in
do f ← add_stmt env label_env (f_body f) [ ] emptyfn;
do u1 ← check_universe_ok ? (f_labgen f);
do u2 ← check_universe_ok ? (f_reggen f);
OK ? f
.
>lookup_add_hit % #E destruct
qed.

definition cminor_noinit_to_rtlabs : Cminor_noinit_program → res RTLabs_program ≝
λp.transform_partial_program … p (transf_partial_fundef … c2ra_function).

include "Cminor/initialisation.ma".

definition cminor_to_rtlabs : Cminor_program → res RTLabs_program ≝
λp. let p' ≝ replace_init p in cminor_noinit_to_rtlabs p'.