include "basics/lists/listb.ma".

include "Clight/Csyntax.ma".

(* TODO: - check signedness - it's almost certainly wrong
         - this is horrific, let's do it further down the compilation chain
           and tweak the semantics to issue extra cost labels for particular
           operations. *)

(* We define replacements for operations on a particular size of integers. *)

definition specific_op ≝ (unary_operation ⊎ binary_operation) × intsize.

definition unop_dec_eq : ∀x,y:unary_operation. (x = y) ⊎ (x ≠ y).
* *; try @(inl … (refl …)) %2 % #E destruct
qed.

definition binop_dec_eq : ∀x,y:binary_operation. (x = y) ⊎ (x ≠ y).
* *; try @(inl … (refl …)) %2 % #E destruct
qed.

definition specific_op_dec_eq : ∀x,y:specific_op. (x = y) ⊎ (x ≠ y).
* * #op #sz * * #op' #sz'
[ 2,3: %2 % #E destruct ]
@(eq_intsize_elim sz sz')
[ 1,3: * #NE %2 % #E destruct @NE @refl ]
#E >E
[ cases (unop_dec_eq op op')
  [ #E2 >E2 %1 @refl
  | * #NE %2 % #E2 @NE destruct @refl
  ]
| cases (binop_dec_eq op op')
  [ #E2 >E2 %1 @refl
  | * #NE %2 % #E2 @NE destruct @refl
  ]
] qed.

(* How to specify a replacement *)

record replace_op : Type[0] ≝
{ op_to_replace : specific_op
; op_requires : list specific_op
; function_id : ident
; function_fd : clight_fundef
}.

(* Find out which operations are used so that we only add the runtime support
   functions that we need. *)

definition add_op ≝
λop:unary_operation ⊎ binary_operation.λty,tl.
  match ty with
  [ Tint sz _ ⇒ 〈op,sz〉::tl
  | _ ⇒ tl
  ].

let rec expr_ops (e:expr) : list specific_op ≝
match e with
[ Expr e _ ⇒
  match e with
  [ Ederef e  ⇒ expr_ops e
  | Eaddrof e ⇒ expr_ops e
  | Eunop op e ⇒ add_op (inl … op) (typeof e) (expr_ops e)
  | Ebinop op e1 e2 ⇒ add_op (inr … op) (typeof e1) (expr_ops e1 @ expr_ops e2)
  | Ecast _ e ⇒ expr_ops e
  | Econdition e1 e2 e3 ⇒ expr_ops e1 @ expr_ops e2 @ expr_ops e3
  | Eandbool e1 e2 ⇒ expr_ops e1 @ expr_ops e2
  | Eorbool e1 e2 ⇒ expr_ops e1 @ expr_ops e2
  | Efield e _ ⇒ expr_ops e
  | Ecost _ e  ⇒ expr_ops e
  | _ ⇒ [ ]
  ]
].

let rec stmt_ops (s:statement) : list specific_op ≝
match s with
[ Sassign e1 e2 ⇒ expr_ops e1 @ expr_ops e2
| Scall oe ef es ⇒ let l ≝ foldr ?? (λe,l. expr_ops e @ l) (expr_ops ef) es in
                   match oe with [ Some e ⇒ expr_ops e @ l | None ⇒ l ]
| Ssequence s1 s2 ⇒ stmt_ops s1 @ stmt_ops s2
| Sifthenelse e1 s1 s2 ⇒ expr_ops e1 @ stmt_ops s1 @ stmt_ops s2
| Swhile e1 s1   ⇒ expr_ops e1 @ stmt_ops s1
| Sdowhile e1 s1 ⇒ expr_ops e1 @ stmt_ops s1
| Sfor s1 e1 s2 s3 ⇒ expr_ops e1 @ stmt_ops s1 @ stmt_ops s2 @ stmt_ops s3
| Sreturn oe ⇒ match oe with [ Some e ⇒ expr_ops e | None ⇒ [ ] ]
| Sswitch e1 ls ⇒ expr_ops e1 @ lstmt_ops ls
| Slabel _ s1 ⇒ stmt_ops s1
| Scost _ s1 ⇒ stmt_ops s1
| _ ⇒ [ ]
]
and lstmt_ops (ls:labeled_statements) : list ((unary_operation ⊎ binary_operation) × intsize) ≝
match ls with
[ LSdefault s1 ⇒ stmt_ops s1
| LScase _ _ s1 ls ⇒ stmt_ops s1 @ lstmt_ops ls
].

definition prog_ops : clight_program → list specific_op ≝
λp. foldr ?? (λf,l. match \snd f with [ CL_Internal f ⇒ stmt_ops (fn_body f) @ l | _ ⇒ l ]) [ ] (prog_funct ?? p).

(* Actually add the functions. *)

definition add_op_replacements : clight_program → list replace_op → clight_program ≝
λp,rs.
  let ops ≝ prog_ops p in
  let 〈ops,extra〉 ≝ foldr ?? (λr,a.
    let 〈ops,fns〉 ≝ a in
    if exists ? (λo. match specific_op_dec_eq o (op_to_replace r) with [ inl _ ⇒ true | _ ⇒ false ]) ops
    then 〈op_requires r @ ops, 〈function_id r, function_fd r〉::fns〉 else 〈ops,fns〉) 〈ops, [ ]〉 rs in
  mk_program ?? (prog_vars ?? p) (extra @ prog_funct ?? p) (prog_main ?? p).


(* Now replace the operations by function calls.  Unfortunately, Clight is
   about the worst possible place to do this - we need to split up expressions
   because we can't put a function call in an expression.
   
   Also, we need to thread name generation throughout. *)

definition find_replacement ≝
λop,ty,rs.
  match ty with
  [ Tint sz _ ⇒ find ?? (λr.match specific_op_dec_eq 〈op,sz〉 (op_to_replace r) with [ inl _ ⇒ Some ? 〈function_id r, type_of_fundef (function_fd r)〉 | _ ⇒ None ? ]) rs
  | _ ⇒ None ?
  ].

let rec rev_prepend_stmts (l:list statement) (s:statement) : statement ≝
match l with
[ nil ⇒ s
| cons h t ⇒ rev_prepend_stmts t (Ssequence h s)
].

let rec subst_expr (e:expr) (rs:list replace_op) (u:universe SymbolTag × (list (ident×type))) : (list statement) × expr × (universe SymbolTag × (list (ident×type))) ≝
match e with
[ Expr e ty ⇒
  match e with
  [ Ederef e1 ⇒ let 〈ls,e1,u〉 ≝ subst_expr e1 rs u in 〈ls, Expr (Ederef e1) ty, u〉
  | Eaddrof e1 ⇒ let 〈ls,e1,u〉 ≝ subst_expr e1 rs u in 〈ls, Expr (Eaddrof e1) ty, u〉
  | Eunop op e1 ⇒
      let 〈ls,e1,u〉 ≝ subst_expr e1 rs u in
      match find_replacement (inl ?? op) (typeof e1) rs with
      [ Some idty ⇒ 
          let 〈tmp,u1〉 ≝ fresh … (\fst u) in
          〈(Scall (Some ? (Expr (Evar tmp) ty)) (Expr (Evar (\fst idty)) (\snd idty)) [e1])::ls,
           (Expr (Evar tmp) ty), 〈u1, 〈tmp, ty〉::(\snd u)〉 〉
      | None ⇒ 〈ls, (Expr (Eunop op e1) ty), u〉
      ]
  | Ebinop op e1 e2 ⇒
      let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
      let 〈ls2,e2,u〉 ≝ subst_expr e2 rs u in
      match find_replacement (inr ?? op) (typeof e1) rs with
      [ Some idty ⇒ 
          let 〈tmp,u1〉 ≝ fresh … (\fst u) in
          〈(Scall (Some ? (Expr (Evar tmp) ty)) (Expr (Evar (\fst idty)) (\snd idty)) [e1; e2])::ls1@ls2,
           (Expr (Evar tmp) ty), 〈u1, 〈tmp, ty〉::(\snd u)〉〉
      | None ⇒ 〈ls1@ls2, (Expr (Ebinop op e1 e2) ty), u〉
      ]
  | Ecast ty1 e1 ⇒ let 〈ls,e1,u〉 ≝ subst_expr e1 rs u in 〈ls, Expr (Ecast ty1 e1) ty, u〉
  | Econdition e1 e2 e3 ⇒
      let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
      let 〈ls2,e2,u〉 ≝ subst_expr e2 rs u in
      let 〈ls3,e3,u〉 ≝ subst_expr e3 rs u in
        〈ls1@ls2@ls3, Expr (Econdition e1 e2 e3) ty, u〉
  | Eandbool e1 e2 ⇒
      let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
      let 〈ls2,e2,u〉 ≝ subst_expr e2 rs u in
        〈ls1@ls2, Expr (Eandbool e1 e2) ty, u〉
  | Eorbool e1 e2 ⇒
      let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
      let 〈ls2,e2,u〉 ≝ subst_expr e2 rs u in
        〈ls1@ls2, Expr (Eorbool e1 e2) ty, u〉
  | Efield e1 id ⇒ let 〈ls,e1,u〉 ≝ subst_expr e1 rs u in 〈ls, Expr (Efield e1 id) ty, u〉
  | Ecost l e1 ⇒ let 〈ls,e1,u〉 ≝ subst_expr e1 rs u in 〈ls, Expr (Ecost l e1) ty, u〉
  | _ ⇒ 〈[ ], Expr e ty, u〉
  ]
].

let rec append_stmts (s:statement) (ls:list statement) on ls : statement ≝
match ls with
[ nil ⇒ s
| cons h t ⇒ append_stmts (Ssequence s h) t
].

let rec subst_stmt (s:statement) (rs:list replace_op) (u:universe SymbolTag × (list (ident×type))) : statement × (universe SymbolTag × (list (ident×type))) ≝
match s with
[ Sassign e1 e2 ⇒
    let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
    let 〈ls2,e2,u〉 ≝ subst_expr e2 rs u in
    〈rev_prepend_stmts (ls1@ls2) (Sassign e1 e2), u〉
| Scall oe ef es ⇒
    let 〈ls1,oe,u〉 ≝
      match oe with
      [ Some e ⇒ let 〈ls1,e,u〉 ≝ subst_expr e rs u in 〈ls1, Some ? e, u〉
      | None ⇒ 〈[ ], None ?, u〉
      ] in
    let 〈ls2,ef,u〉 ≝ subst_expr ef rs u in
    let 〈lss,es,u〉 ≝ foldr ?? (λe,x.
      let 〈lss,es,u〉 ≝ x in 
      let 〈ls,e,u〉 ≝ subst_expr e rs u in
      〈ls@lss,e::es,u〉) 〈[ ], [ ], u〉 es
    in 〈rev_prepend_stmts (ls1@ls2@lss) (Scall oe ef es), u〉
| Ssequence s1 s2 ⇒
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    let 〈s2,u〉 ≝ subst_stmt s2 rs u in
      〈Ssequence s1 s2, u〉
| Sifthenelse e1 s1 s2 ⇒
    let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    let 〈s2,u〉 ≝ subst_stmt s2 rs u in
      〈rev_prepend_stmts ls1 (Sifthenelse e1 s1 s2), u〉

  (* This is particularly bad - for loops we end up duplicating the code. *)
  
| Swhile e1 s1 ⇒
    let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
      〈rev_prepend_stmts ls1 (Swhile e1 (append_stmts s1 ls1)), u〉
| Sdowhile e1 s1 ⇒
    let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
      〈Sdowhile e1 (append_stmts s1 ls1), u〉
| Sfor s1 e1 s2 s3 ⇒
    let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    let 〈s2,u〉 ≝ subst_stmt s2 rs u in
    let 〈s3,u〉 ≝ subst_stmt s3 rs u in
      〈Sfor (append_stmts s1 ls1) e1 s2 (append_stmts s3 ls1), u〉
| Sreturn oe ⇒
    match oe with
    [ Some e ⇒ let 〈ls1,e,u〉 ≝ subst_expr e rs u in 〈rev_prepend_stmts ls1 (Sreturn (Some ? e)), u〉
    | None ⇒ 〈Sreturn (None ?), u〉
    ]
| Sswitch e1 lbs ⇒
    let 〈ls1,e1,u〉 ≝ subst_expr e1 rs u in
    let 〈lbs,u〉 ≝ subst_lstmts lbs rs u in
      〈rev_prepend_stmts ls1 (Sswitch e1 lbs), u〉
| Slabel l s1 ⇒
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    〈Slabel l s1, u〉
| Scost l s1 ⇒
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    〈Scost l s1, u〉
| _ ⇒ 〈s, u〉
]
and subst_lstmts (lbs:labeled_statements) rs u ≝
match lbs with
[ LSdefault s1 ⇒
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    〈LSdefault s1, u〉
| LScase sz i s1 lbs ⇒
    let 〈s1,u〉 ≝ subst_stmt s1 rs u in
    let 〈lbs,u〉 ≝ subst_lstmts lbs rs u in
    〈LScase sz i s1 lbs, u〉
].

definition subst_fn : function → list replace_op → universe SymbolTag → function × (universe SymbolTag) ≝
λf,rs,u.
  let 〈body,u〉 ≝ subst_stmt (fn_body f) rs 〈u, [ ]〉 in
  let 〈u,vs〉 ≝ u in
    〈mk_function (fn_return f) (fn_params f) (vs @ (fn_vars f)) body, u〉.

definition subst_fundef : clight_fundef → list replace_op → universe SymbolTag → clight_fundef × (universe SymbolTag) ≝
λfd,rs,u.
  match fd with
  [ CL_Internal f ⇒ let 〈f,u〉 ≝ subst_fn f rs u in 〈CL_Internal f, u〉
  | _ ⇒ 〈fd, u〉
  ].

definition subst_prog : clight_program → list replace_op → universe SymbolTag → clight_program × (universe SymbolTag) ≝
λp,rs,u.
  let 〈fns, u〉 ≝ foldr ?? (λf,a.
    let 〈fns, u〉 ≝ a in
    let 〈fd,u〉 ≝ subst_fundef (\snd f) rs u in
    〈〈\fst f, fd〉::fns, u〉) 〈[ ], u〉 (prog_funct ?? p)
  in 〈mk_program ?? (prog_vars ?? p) fns (prog_main ?? p), u〉.

(* Now define some functions in Clight. *)

definition divu : universe SymbolTag → intsize → replace_op × (universe SymbolTag) ≝
λu,sz.
  let uint ≝ Tint sz Unsigned in
  let 〈fn_id, u〉 ≝ fresh … u in
  let 〈x, u〉 ≝ fresh … u in
  let 〈y, u〉 ≝ fresh … u in
  let 〈quo, u〉 ≝ fresh … u in
  〈mk_replace_op 〈inr ?? Odiv, sz〉 [ ] fn_id
    (CL_Internal (mk_function uint
      [〈x, uint〉; 〈y, uint〉]
      [〈quo, uint〉]
      (Ssequence
        (Sassign (Expr (Evar quo) uint) (Expr (Econst_int sz (repr sz 0)) uint))
      (Ssequence
        (Swhile
          (Expr (Ebinop Oge (Expr (Evar x) uint) (Expr (Evar y) uint)) uint)
          (Ssequence
            (Sassign (Expr (Evar x) uint) (Expr (Ebinop Osub (Expr (Evar x) uint) (Expr (Evar y) uint)) uint))
            (Sassign (Expr (Evar quo) uint) (Expr (Ebinop Oadd (Expr (Evar quo) uint) (Expr (Econst_int sz (repr sz 1)) uint)) uint))))
        (Sreturn (Some ? (Expr (Evar quo) uint))))))),
    u〉.
        
(* And do everything... *)

definition add_runtime_replacements : clight_program → universe SymbolTag → clight_program × (universe SymbolTag) ≝
λp,u.
  let 〈divu16,u〉 ≝ divu u I16 in
  let 〈divu32,u〉 ≝ divu u I32 in
  let rs ≝ [divu16;divu32] in
  let p ≝ add_op_replacements p rs in
  subst_prog p rs u.