source: Deliverables/D2.2/8051/src/clight/loopUnrolling.ml @ 1569

open Clight

module IdentSet = Set.Make(struct type t = AST.ident let compare = compare end)

type heuristic_global_info = int * IdentSet.t

type heuristic_local_info = IdentSet.t

let big_union : IdentSet.t list -> IdentSet.t =
  List.fold_left IdentSet.union IdentSet.empty

let referenced : statement -> IdentSet.t =
  let f_expr e res_list =
    let res = big_union res_list in
    let Expr(ed, _) = e in
    match ed with
      | Eaddrof (Expr(Evar x, _)) -> IdentSet.add x res
      | _ -> res in
  let f_stmt stmt eres sres =
    IdentSet.union (big_union eres) (big_union sres) in
  ClightFold.statement2 f_expr f_stmt

let mk_heuristic_global_info (max_factor : int) (p : program)
    : heuristic_global_info =
  let f set ((id, _), _) = IdentSet.add id set in
  (max_factor, List.fold_left f IdentSet.empty p.prog_vars)

let mk_heuristic_local_info (f : cfunction) =   referenced f.fn_body

let vars_of_expr, vars_of_stmt =
  let f_expr e res_list =
    let res = big_union res_list in
    let Expr(ed, _) = e in
    match ed with
      | Evar x -> IdentSet.add x res
      | _ -> res in
  let f_stmt stmt eres sres =
    IdentSet.union (big_union eres) (big_union sres) in
  ClightFold.expr2 f_expr, ClightFold.statement2 f_expr f_stmt

let changed_in =
  let f_expr e _ = () in
  let f_stmt stmt _ sres =
    let res = big_union sres in
    match stmt with
      | Sassign (Expr(Evar x, _), _)
      | Clight.Scall (Some (Expr(Evar x, _)), _, _) ->
        (* a situation like "*&x = ... " is avoided as dereferenced variables*)
        (* will be excluded a priori *)
        IdentSet.add x res
      | _ -> res in
  ClightFold.statement2 f_expr f_stmt

type sign =
  | Plus
  | Minus

let find_increment = function
  | Sassign(
    Expr(Evar x,_),
    Expr(
      (Ebinop(Oadd,Expr(Evar y,_),inc) |
          Ebinop(Oadd,inc,Expr(Evar y,_)) ),_))
      when x = y -> Some (x, Plus, inc)
  | Sassign(
    Expr(Evar x,_),
    Expr(
      (Ebinop(Osub,Expr(Evar y,_),inc) |
          Ebinop(Osub,inc,Expr(Evar y,_)) ),_))
      when x = y -> Some (x, Minus, inc)
  | _ -> None

let find_bound s x (Expr(ed, _)) = match s, ed with
  | Plus, Ebinop(Olt,Expr(Evar y,_),b)
  | Plus, Ebinop(Ogt,b,Expr(Evar y,_))
  | Minus, Ebinop(Olt,b,Expr(Evar y,_))
  | Minus, Ebinop(Ogt,Expr(Evar y,_),b)
    when y = x -> Some (false, b)
  | Plus, Ebinop(Ole,Expr(Evar y,_),b)
  | Plus, Ebinop(Oge,b,Expr(Evar y,_))
  | Minus, Ebinop(Ole,b,Expr(Evar y,_))
  | Minus, Ebinop(Oge,Expr(Evar y,_),b)
    when y = x -> Some (true, b)
  | _ -> None

(* this function tells when to perform the optimization on a statement. It *)
(* should always return [false] when the optimization cannot be applied *)
let heuristics
    ((max_factor, global_vars) : heuristic_global_info)
    (rfrncd_vars : heuristic_local_info)
    : Clight.statement -> (int * AST.ident * sign * expr * expr * bool) option =
  let global_referenced = IdentSet.union global_vars rfrncd_vars in
  function
    | Clight.Sfor (Some _, init, guard, tail, body) ->
      begin
        match find_increment tail with
          | Some (x, dir, inc) ->
            begin
              match find_bound dir x guard with
                | Some (included, bound) ->
                  let head_vars =
                    IdentSet.union (vars_of_expr guard) (vars_of_stmt tail) in
                  let dangerous =
                    IdentSet.union (changed_in body) global_referenced in
                  if IdentSet.is_empty (IdentSet.inter head_vars dangerous) then
                    Some (max_factor, x, dir, inc, bound, included)
                  else
                    None
                | None -> None
            end
          | None -> None
      end
    (* TODO add a check for size *)
    | _ -> None

let sexpr a b = CostLabel.Sexpr(a,b)

(** compose the mapping i |--> s with all labels *)
let reindex_expr, reindex_stmt =
  let f_expr i s e expr_res =
    match e, expr_res with
      | Clight.Expr (Clight.Ecost (lbl, _), t), e::_ ->
        let lbl = CostLabel.comp_index i s lbl in
        Clight.Expr (Clight.Ecost (lbl, e), t)
      | _ -> ClightFold.expr_fill_exprs e expr_res in
  let f_stmt i s stmt expr_res stmt_res =
    match stmt, stmt_res with
      | Clight.Scost(lbl, _), stmt' :: _ ->
        let lbl = CostLabel.comp_index i s lbl in
        Clight.Scost(lbl,stmt')
      | _ -> ClightFold.statement_fill_subs stmt expr_res stmt_res in
  (fun i s -> ClightFold.expr2 (f_expr i s)),
  (fun i s -> ClightFold.statement2 (f_expr i s) (f_stmt i s))

let labels_of =
  let f_expr _ _ = () in
  let f_stmt stmt _ res_stmts =
    let s = match stmt with
      | Clight.Slabel(lbl, _) -> Label.Set.singleton lbl
      | _ -> Label.Set.empty in
    List.fold_left Label.Set.union s res_stmts in
  ClightFold.statement2 f_expr f_stmt

let create_fresh_labels fresh lbl_set =
  let f_lbl lbl = Label.Map.add lbl (fresh ()) in
  Label.Set.fold f_lbl lbl_set Label.Map.empty

let apply_label_map map =
  let f_stmt stmt exprs stmts =
    match stmt, stmts with
      | Clight.Slabel (lbl, _), s :: _ ->
                                (* lbl must be in domain of map *)
        Clight.Slabel (Label.Map.find lbl map, s)
      | Clight.Sgoto lbl, _ ->
        Clight.Sgoto (
          try
            Label.Map.find lbl map
          with
            | Not_found -> (* means it is an outgoing goto *)
              lbl)
      | _ -> ClightFold.statement_fill_subs stmt exprs stmts in
  ClightFold.statement2 ClightFold.expr_fill_exprs f_stmt

(* unroll the k-th copy over a total of n *)
let rec do_the_unrolling fresh k n i e sb st labels_of_sb =
  if k = 0 then reindex_stmt i (sexpr n 0) sb else
    let preamble = do_the_unrolling fresh (k-1) n i e sb st labels_of_sb in
    let st_pre = reindex_stmt i (sexpr n (k-1)) st in
    let e_next = reindex_expr i (sexpr n k) e in
    let sb_next = apply_label_map (create_fresh_labels fresh labels_of_sb) sb in
    let sb_next = reindex_stmt i (sexpr n k) sb_next in
    Ssequence(preamble,
              Ssequence(st_pre,
                        Ssequence(Sifthenelse(e_next, Sskip, Sbreak),
                                  sb_next)))

let unroll fresh g_info l_info =
  let f_stmt stmt exprs stmts =
    match heuristics g_info l_info stmt, stmt, exprs, stmts with
      | Some (k, x, dir, inc, bound, included),
    Clight.Sfor (Some i, _, _, _, _), e :: _, si :: st :: sb ::_ ->
      (* we can suppose no label in s3 is target of a goto outside of stmt, *)
      (* as loop is indexed and thus single entry. So we can freely rename *)
      (* labels in stmt. *)
      (* IMPORTANT: I am supposing no labels are in s1 and s2 *)
      let labels_of_sb = labels_of sb in
      let e_first = reindex_expr i (sexpr k 0) e in
      let st_last = reindex_stmt i (sexpr k (k-1)) st in
      let body = do_the_unrolling fresh (k-1) k i e sb st labels_of_sb in
      Sfor (Some i, si, e_first, st_last, body)
      | Some _, _, _, _ ->
        assert false (* heuristics should have crossed out other cases *)
      | _ -> ClightFold.statement_fill_subs stmt exprs stmts in
  ClightFold.statement2 ClightFold.expr_fill_exprs f_stmt

let unroll_funct fresh g_info = function
  | (id, Clight.Internal def) ->
    let l_info = mk_heuristic_local_info def in
    let body = unroll fresh g_info l_info def.Clight.fn_body in
    (id, Clight.Internal {def with Clight.fn_body = body})
  | f -> f

open Languages

let trans ?(factor = 2) () = Clight, function
  | AstClight p ->
    let g_info = mk_heuristic_global_info factor p in
    let fresh = ClightAnnotator.make_fresh "_l" p in
    let functs = List.map (unroll_funct fresh g_info) p.Clight.prog_funct in
    AstClight { p with Clight.prog_funct = functs }
  | _ -> assert false