include "Clight/Csyntax.ma".

(* Extract cost labels from a program. *)

let rec labels_of_expr (e:expr) : list costlabel ≝
match e with
[ Expr e' _ ⇒
  match e' with
  [ Ederef e1 ⇒ labels_of_expr e1
  | Eaddrof e1 ⇒ labels_of_expr e1
  | Eunop _ e1 ⇒ labels_of_expr e1
  | Ebinop _ e1 e2 ⇒ labels_of_expr e1 @ labels_of_expr e2
  | Ecast _ e1 ⇒ labels_of_expr e1
  | Econdition e1 e2 e3 ⇒ labels_of_expr e1 @ labels_of_expr e2 @ labels_of_expr e3
  | Eandbool e1 e2 ⇒ labels_of_expr e1 @ labels_of_expr e2
  | Eorbool e1 e2 ⇒ labels_of_expr e1 @ labels_of_expr e2
  | Efield e1 _ ⇒ labels_of_expr e1
  | Ecost cl e1 ⇒ cl::(labels_of_expr e1)
  | _ ⇒ [ ]
  ]
].

let rec labels_of_statement (s:statement) : list costlabel ≝
match s with
[ Sassign e1 e2 ⇒ labels_of_expr e1 @ labels_of_expr e2
| Scall oe e1 es ⇒ option_map_def … labels_of_expr [ ] oe @ labels_of_expr e1 @ foldl … (λls,e. labels_of_expr e @ ls) [ ] es
| Ssequence s1 s2 ⇒ labels_of_statement s1 @ labels_of_statement s2
| Sifthenelse e1 s1 s2 ⇒ labels_of_expr e1 @ labels_of_statement s1 @ labels_of_statement s2
| Swhile e1 s1 ⇒ labels_of_expr e1 @ labels_of_statement s1
| Sdowhile e1 s1 ⇒ labels_of_expr e1 @ labels_of_statement s1
| Sfor s1 e1 s2 s3 ⇒ labels_of_statement s1 @ labels_of_expr e1 @ labels_of_statement s2 @ labels_of_statement s3
| Sreturn oe ⇒ option_map_def … labels_of_expr [ ] oe
| Sswitch e1 ls ⇒ labels_of_expr e1 @ labels_of_labeled_statements ls
| Slabel _ s1 ⇒ labels_of_statement s1
| Scost cl s1 ⇒ cl::(labels_of_statement s1)
| _ ⇒ [ ]
] and labels_of_labeled_statements (ls:labeled_statements) : list costlabel ≝
match ls with
[ LSdefault s1 ⇒ labels_of_statement s1
| LScase _ _ s1 ls1 ⇒ labels_of_statement s1 @ labels_of_labeled_statements ls1
].

definition labels_of_clight_fundef : ident × clight_fundef → list costlabel ≝
λifd. match \snd ifd with
[ CL_Internal f ⇒ labels_of_statement (fn_body f)
| _ ⇒ [ ]
].

definition labels_of_clight : clight_program → list costlabel ≝
λp. foldl … (λls,f. labels_of_clight_fundef f @ ls) [ ] (prog_funct ?? p).

definition in_clight_program : clight_program → costlabel → Prop ≝
λp,l. Exists … (λx.x=l) (labels_of_clight p).

definition in_clight_label ≝
  λp. Σl. in_clight_program p l. 

definition clight_cost_map ≝
  λp. (in_clight_label p) → ℕ.

definition clight_label_free : clight_program → bool ≝
λp. match labels_of_clight p with [ nil ⇒ true | _ ⇒ false ].

(* Adding labels *)

definition add_cost_before : statement → universe CostTag → statement × (universe CostTag) ≝
λs,gen.
  let 〈l,gen〉 ≝ fresh … gen in
  〈Scost l s, gen〉.

definition add_cost_after : statement → universe CostTag → statement × (universe CostTag) ≝
λs,gen.
  let 〈l,gen〉 ≝ fresh … gen in
  〈Ssequence s (Scost l Sskip), gen〉.

definition add_cost_expr : expr → universe CostTag → expr × (universe CostTag) ≝
λe,gen.
  let 〈l,gen〉 ≝ fresh … gen in
  〈Expr (Ecost l e) (typeof e), gen〉.

definition const_int : intsize → nat → expr ≝
λsz,n. Expr (Econst_int ? (repr sz n)) (Tint sz Signed).

let rec label_expr (e:expr) (costgen:universe CostTag)
        on e : expr × (universe CostTag) ≝
match e with
[ Expr ed ty ⇒ let 〈ed,costgen〉 ≝ label_expr_descr ed costgen ty in
               〈Expr ed ty, costgen〉
]
and label_expr_descr (e:expr_descr) (costgen:universe CostTag) (ty:type)
    on e : expr_descr × (universe CostTag) ≝
match e with
[ Ederef e' ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    〈Ederef e', costgen〉
| Eaddrof e' ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    〈Eaddrof e', costgen〉
| Eunop op e' ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    〈Eunop op e', costgen〉
| Ebinop op e1 e2 ⇒
    let 〈e1,costgen〉 ≝ label_expr e1 costgen in 
    let 〈e2,costgen〉 ≝ label_expr e2 costgen in 
    〈Ebinop op e1 e2, costgen〉
| Ecast ty e' ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    〈Ecast ty e', costgen〉
| Econdition e' e1 e2 ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    let 〈e1,costgen〉 ≝ label_expr e1 costgen in 
    let 〈e1,costgen〉 ≝ add_cost_expr e1 costgen in 
    let 〈e2,costgen〉 ≝ label_expr e2 costgen in 
    let 〈e2,costgen〉 ≝ add_cost_expr e2 costgen in 
    〈Econdition e' e1 e2, costgen〉
(* andbool and orbool are changed to conditionals to capture their
   short-circuiting cost difference; note that we have to return 0 or 1, and
   we get rather more cost labels than I'd like *)
| Eandbool e1 e2 ⇒
    let 〈e1,costgen〉 ≝ label_expr e1 costgen in 
    let 〈e2,costgen〉 ≝ label_expr e2 costgen in
    match ty with
    [ Tint sz sg ⇒
      let 〈et,costgen〉 ≝ add_cost_expr (const_int sz 1) costgen in
      let 〈ef,costgen〉 ≝ add_cost_expr (const_int sz 0) costgen in
      let 〈e2,costgen〉 ≝ add_cost_expr (Expr (Econdition e2 et ef) ty) costgen in 
      let 〈ef,costgen〉 ≝ add_cost_expr (const_int sz 0) costgen in
      〈Econdition e1 e2 ef, costgen〉
    | _ ⇒ (* default on 32 bit consts if inconsistent type. *)
      let 〈et,costgen〉 ≝ add_cost_expr (const_int I32 1) costgen in
      let 〈ef,costgen〉 ≝ add_cost_expr (const_int I32 0) costgen in
      let 〈e2,costgen〉 ≝ add_cost_expr (Expr (Econdition e2 et ef) ty) costgen in 
      let 〈ef,costgen〉 ≝ add_cost_expr (const_int I32 0) costgen in
      〈Econdition e1 e2 ef, costgen〉        
    ]
| Eorbool e1 e2 ⇒
    let 〈e1,costgen〉 ≝ label_expr e1 costgen in 
    let 〈e2,costgen〉 ≝ label_expr e2 costgen in 
    match ty with
    [ Tint sz sg ⇒    
      let 〈et,costgen〉 ≝ add_cost_expr (const_int sz 1) costgen in
      let 〈ef,costgen〉 ≝ add_cost_expr (const_int sz 0) costgen in
      let 〈e2,costgen〉 ≝ add_cost_expr (Expr (Econdition e2 et ef) ty) costgen in 
      let 〈et,costgen〉 ≝ add_cost_expr (const_int sz 1) costgen in
      〈Econdition e1 et e2, costgen〉
    | _ ⇒
      let 〈et,costgen〉 ≝ add_cost_expr (const_int I32 1) costgen in
      let 〈ef,costgen〉 ≝ add_cost_expr (const_int I32 0) costgen in
      let 〈e2,costgen〉 ≝ add_cost_expr (Expr (Econdition e2 et ef) ty) costgen in 
      let 〈et,costgen〉 ≝ add_cost_expr (const_int I32 1) costgen in
      〈Econdition e1 et e2, costgen〉
    ]
| Efield e' id ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    〈Efield e' id, costgen〉
(* The prototype asserts on preexisting cost labels, but I'd quite like to
   keep them. *)
| Ecost l e' ⇒
    let 〈e',costgen〉 ≝ label_expr e' costgen in 
    〈Ecost l e', costgen〉

(* The remaining expressions don't have subexpressions. *)
| _ ⇒ 〈e,costgen〉
].

let rec label_exprs (l:list expr) (costgen:universe CostTag)
        on l : list expr × (universe CostTag) ≝
match l with
[ nil ⇒ 〈nil ?,costgen〉
| cons e es ⇒
    let 〈e,costgen〉 ≝ label_expr e costgen in
    let 〈es,costgen〉 ≝ label_exprs es costgen in
    〈e::es,costgen〉
].

definition label_opt_expr ≝
λoe,costgen. match oe with
[ None ⇒ 〈None ?, costgen〉
| Some e ⇒ let 〈e,costgen〉 ≝ label_expr e costgen in 〈Some ? e,costgen〉
].


let rec label_statement (s:statement) (costgen:universe CostTag)
        on s : statement × (universe CostTag) ≝
match s with
[ Sskip ⇒ 〈Sskip,costgen〉
| Sassign e1 e2 ⇒
    let 〈e1,costgen〉 ≝ label_expr e1 costgen in 
    let 〈e2,costgen〉 ≝ label_expr e2 costgen in
    〈Sassign e1 e2, costgen〉
| Scall e_ret e_fn e_args ⇒
    let 〈e_ret,costgen〉 ≝ label_opt_expr e_ret costgen in
    let 〈e_fn,costgen〉 ≝ label_expr e_fn costgen in
    let 〈e_args,costgen〉 ≝ label_exprs e_args costgen in
    〈Scall e_ret e_fn e_args, costgen〉
| Ssequence s1 s2 ⇒
    let 〈s1,costgen〉 ≝ label_statement s1 costgen in
    let 〈s2,costgen〉 ≝ label_statement s2 costgen in
    〈Ssequence s1 s2, costgen〉
| Sifthenelse e s1 s2 ⇒
    let 〈e,costgen〉 ≝ label_expr e costgen in 
    let 〈s1,costgen〉 ≝ label_statement s1 costgen in
    let 〈s1,costgen〉 ≝ add_cost_before s1 costgen in
    let 〈s2,costgen〉 ≝ label_statement s2 costgen in
    let 〈s2,costgen〉 ≝ add_cost_before s2 costgen in
    〈Sifthenelse e s1 s2, costgen〉
| Swhile e s' ⇒
    let 〈e,costgen〉 ≝ label_expr e costgen in 
    let 〈s',costgen〉 ≝ label_statement s' costgen in
    let 〈s',costgen〉 ≝ add_cost_before s' costgen in
    let 〈s,costgen〉 ≝ add_cost_after (Swhile e s') costgen in
    〈s,costgen〉
| Sdowhile e s' ⇒
    let 〈e,costgen〉 ≝ label_expr e costgen in 
    let 〈s',costgen〉 ≝ label_statement s' costgen in
    let 〈s',costgen〉 ≝ add_cost_before s' costgen in
    let 〈s,costgen〉 ≝ add_cost_after (Sdowhile e s') costgen in
    〈s,costgen〉
| Sfor s1 e s2 s3 ⇒
    let 〈e,costgen〉 ≝ label_expr e costgen in 
    let 〈s1,costgen〉 ≝ label_statement s1 costgen in
    let 〈s2,costgen〉 ≝ label_statement s2 costgen in
    let 〈s3,costgen〉 ≝ label_statement s3 costgen in
    let 〈s3,costgen〉 ≝ add_cost_before s3 costgen in
    let 〈s,costgen〉 ≝ add_cost_after (Sfor s1 e s2 s3) costgen in
    〈s,costgen〉
| Sbreak ⇒ 〈Sbreak,costgen〉
| Scontinue ⇒ 〈Scontinue,costgen〉
| Sreturn opt_e ⇒
    let 〈opt_e,costgen〉 ≝ label_opt_expr opt_e costgen in
    〈Sreturn opt_e,costgen〉
| Sswitch e ls ⇒
    let 〈e,costgen〉 ≝ label_expr e costgen in 
    let 〈ls,costgen〉 ≝ label_lstatements ls costgen in
    〈Sswitch e ls, costgen〉
| Slabel l s' ⇒
    let 〈s',costgen〉 ≝ label_statement s' costgen in
    let 〈s',costgen〉 ≝ add_cost_before s' costgen in
    〈Slabel l s', costgen〉
| Sgoto l ⇒ 〈Sgoto l, costgen〉

(* The prototype asserts on preexisting cost labels, but I'd quite like to
   keep them. *)
| Scost l s' ⇒
    let 〈s',costgen〉 ≝ label_statement s' costgen in
    〈Scost l s', costgen〉
]
and label_lstatements (ls:labeled_statements) (costgen:universe CostTag)
        on ls : labeled_statements × (universe CostTag) ≝
match ls with
[ LSdefault s ⇒
    let 〈s,costgen〉 ≝ label_statement s costgen in
    let 〈s,costgen〉 ≝ add_cost_before s costgen in
    〈LSdefault s, costgen〉
| LScase sz i s ls' ⇒
    let 〈s,costgen〉 ≝ label_statement s costgen in
    let 〈s,costgen〉 ≝ add_cost_before s costgen in
    let 〈ls',costgen〉 ≝ label_lstatements ls' costgen in
    〈LScase sz i s ls', costgen〉
].

definition label_function : universe CostTag → function → function × (universe CostTag) ≝
λcostgen,f.
  let 〈body,costgen〉 ≝ label_statement (fn_body f) costgen in
  let 〈body,costgen〉 ≝ add_cost_before body costgen in
    〈mk_function (fn_return f) (fn_params f) (fn_vars f) body, costgen〉.

definition label_fundef : universe CostTag → clight_fundef → clight_fundef × (universe CostTag) ≝
λgen,f. match f with
[ CL_Internal f ⇒
    let 〈f',gen'〉 ≝ label_function gen f in
    〈CL_Internal f', gen'〉
| CL_External id args ty ⇒ 〈CL_External id args ty, gen〉
].



definition clight_label : clight_program → clight_program × costlabel ≝
λp.
  let costgen ≝ new_universe CostTag in
  let 〈init_cost, costgen〉 ≝ fresh … costgen in
  〈\fst (transform_program_gen … costgen p (λ_.label_fundef)), init_cost〉.