source: Deliverables/D2.2/8051/src/common/costExpr.ml

type cond =
        | Ceq of int         (** index is equal to *)
        | Cgeq of int        (** index is greater or equal to *)
        | Cmod of int*int      (** index modulo equal to *)
  | Cgeqmod of int*int*int (** index greater than and modulo equal to *)

module CondSet = Set.Make(struct
    type t = cond
    let compare = compare
end)

open CostLabel

    
let cond_of_sexpr = function
  | Sexpr(0, b) -> Ceq(b)
  | Sexpr(1, b) -> Cgeq(b)
  | Sexpr(a, b) when b < a -> Cmod(a, b)
  | Sexpr(a, b) -> Cgeqmod(b, a, b mod a)

type cost_expr =
    | Exact of int
    | Ternary of index * CondSet.t * cost_expr * cost_expr

(* compute from the set [s] a the 3-uple [Some (h, s_h, s_rest)] where *)
(* [h] is the head of first elements of [s], and [s] is the union of *)
(* [{h :: tl | tl in s_h}] and [s_rest]. Gives [None] if either [s] is empty *)
(* or it starts with an empty list. ([s] should contain lists with same length *)
let heads_tails_of s =
    if IndexingSet.is_empty s then None else
    match IndexingSet.min_elt s with
        | [] -> None
        | head :: _ ->
            let filter x = (List.hd x = head) in
            let (s_head, s_rest) = IndexingSet.partition filter s in
            let add_tail l = IndexingSet.add (List.tl l) in
            let s_head = IndexingSet.fold add_tail s_head IndexingSet.empty in 
            Some (head, s_head, s_rest)  

let rec cost_mapping_ind atom ind (m : int Map.t) (s : IndexingSet.t) =
    let lbl = {name = atom; i = ind} in
    match heads_tails_of s with
        | None when Map.mem lbl m-> Exact (Map.find lbl m)
        | None -> Exact 0
        | Some(h, s_head, s_rest) ->
            let i = List.length ind in
            let condition = cond_of_sexpr h in
            let if_true = cost_mapping_ind atom (h :: ind) m s_head in
            let if_false = cost_mapping_ind atom ind m s_rest in
            Ternary(i, CondSet.singleton condition, if_true, if_false)  

let indexing_sets_from_cost_mapping m =
    let f k _ sets =
        let s =
            try
                IndexingSet.add k.i (Atom.Map.find k.name sets)
            with
                | Not_found -> IndexingSet.singleton k.i in
        Atom.Map.add k.name s sets in
    Map.fold f m Atom.Map.empty
    
(* extended_gcd a b = (x, y, gcd(a, b)) where x*a + y*b = gcd(a, b) *)  
let rec extended_gcd a = function
        | 0 -> (1, 0, a)
        | b ->
                let (x, y, r) = extended_gcd b (a mod b) in
          (y, x - (a / b) * y, r)

let opt_bind (t : 'a option) (f : 'a -> 'b option) : 'b option =
        match t with
                | None -> None
                | Some x -> f x 

(* in the following a set of conditions is considered as sequents, i.e. *)
(* {c1, ..., ck} is considered as c1 || ... || ck. The empty set is for false,*)
(* while true is given by Cgeq 0. We will call sets general conditions *)

let print_cond = function
        | Ceq i -> Printf.sprintf "== %d" i
        | Cgeq i -> Printf.sprintf ">= %d" i
        | Cmod (a, b) -> Printf.sprintf "%% %d == %d" a b
        | Cgeqmod (i, a, b) -> Printf.sprintf ">= %d & %% %d == %d" i a b

let print_gen_cond gc =
        let f c = Printf.sprintf "%s || %s" (print_cond c) in
        CondSet.fold f gc "F"

(* cond_and_single c1 c2 gives c1 && c2 as a generalized condition. Recursion*)
(* is only to re-use match cases *)
let cond_and_single c1 c2 = 
        let rec cond_and_single' c1 c2 = match c1, c2 with 
        | Ceq h as c, Ceq k when h = k -> Some c
        | (Ceq h as c), Cgeq k | Cgeq k, (Ceq h as c) when h >= k ->
                Some c
        | (Ceq h as c), Cmod(a, b) | Cmod (a, b), (Ceq h as c) when h mod a = b ->
                Some c
        | (Ceq h as c), Cgeqmod(k, a, b) | Cgeqmod (k, a, b), (Ceq h as c)
           when h mod a = b && h >= k -> Some c
  | Ceq _, _ | _, Ceq _ -> None
  | Cgeq h, Cgeq k -> Some (Cgeq (max h k))
        | Cgeq h, Cmod (a,b) | Cmod (a,b), Cgeq h ->
                if h <= b then Some (Cmod (a, b)) else
                Some (Cgeqmod(h - h mod a + a + b, a, b))
        | Cgeq h, Cgeqmod (k, a, b) | Cgeqmod(k, a, b), Cgeq h ->
    cond_and_single' (Cgeq (max h k)) (Cmod(a, b))
        (* special case of Chinese remainder theorem *)
        | Cmod (a, b), Cmod(c, d) ->
                let (x, y, gcd) = extended_gcd a c in
                if b mod gcd <> d mod gcd then None else
                let a_gcd = a / gcd in
                let lcm = a_gcd * c in
                let res = (b + a_gcd * x * (d - b)) mod lcm in
                Some (Cmod(lcm, res))
        | Cmod (a, b), Cgeqmod(k, c, d) | Cgeqmod(k, c, d), Cmod(a, b) ->
                opt_bind (cond_and_single' (Cmod(a, b)) (Cmod(c,d)))
                (fun x -> cond_and_single' (Cgeq k) x)
        | Cgeqmod (h, a, b), Cgeqmod(k, c, d) ->
                opt_bind (cond_and_single' (Cmod(a, b)) (Cmod(c,d)))
                (fun x -> cond_and_single' (Cgeq (max h k)) x) in
        match cond_and_single' c1 c2 with
                | None -> CondSet.empty
                | Some x -> CondSet.singleton x

(* this generalizes to general conditions for first argument, based on*)
(* (c1 || ... || ck) && c = (c1 && c || ... || ck && c) *)
let cond_and s1 c2 =
        let add_and c1 = CondSet.union (cond_and_single c1 c2) in
        CondSet.fold add_and s1 CondSet.empty

(* this creates the set { f 0, ..., f (n-1) } *)        
let rec init_set f  n =
        if n <= 0 then CondSet.empty else
  CondSet.add (f (n-1)) (init_set f (n-1))

(* cond_and_not_single c1 c2 is equivalent to c1 && !c2 as a generalized *)
(* condition *)
let rec cond_and_not_single c1 c2 =
                match c1, c2 with
                        | Ceq h, Ceq k when h = k -> CondSet.empty
                        | Ceq h, Cgeq k when h >= k -> CondSet.empty
                        | Ceq h, Cmod (a, b) when h mod a = b -> CondSet.empty
                        | Ceq h, Cgeqmod (k, a, b) when h >= k && h mod a = b -> CondSet.empty
            | Ceq _, _ -> CondSet.singleton c1
                        | Cgeq h, Ceq k when k < h -> CondSet.singleton c1
                        | Cgeq h, Ceq k ->
                                (* Ceq h, Ceq (h+1), ... , Ceq (k-1), Cgeq (k+1) *)
                                let s' = init_set (fun x -> Ceq(h + x)) (k - h) in
                                CondSet.add (Cgeq(k+1)) s'
                        | Cgeq h, Cgeq k ->
                                (* if k < h init_set will correctly give an empty set, otherwise *)
                                (* {Ceq h, ... , Ceq (k - 1)} *)
                                init_set (fun x -> Ceq(h + x)) (k - h)
                        | Cmod (a, b), Ceq k when k mod a <> b -> CondSet.singleton c1
                        | Cmod (a, b), Ceq k ->
                                let s' = init_set (fun x -> Ceq(a * x + b)) (k / a) in
                                CondSet.add (Cgeqmod(k + a, a, b)) s'
                        | Cmod (a, b), Cgeq k ->
                                init_set (fun x -> Ceq(a * x + b)) (k / a)
                        | Cgeqmod(h, a, b), Ceq k when k < h || k mod a <> b ->
                                CondSet.singleton c1
                        | Cgeqmod(h, a, b), Ceq k ->
                                let h' = h - h mod a + b in
        let s' = init_set (fun x -> Ceq(a * x + h')) (k / a - h / a) in
                                CondSet.add (Cgeqmod(k + a, a, b)) s'
                        | Cgeqmod(h, a, b), Cgeq k ->
        let h' = h - h mod a + b in
        init_set (fun x -> Ceq(a * x + h')) (k / a - h / a)
                        (* when we do not use cleverer ways, we just use cond_and_single *)
            | c1, (Cmod (a, b) as c2) ->
        let s' = CondSet.remove c2 (init_set (fun x -> Cmod(a, x)) a) in
        let f x = CondSet.union (cond_and_single c1 x) in
        CondSet.fold f s' CondSet.empty
            | c1, Cgeqmod (k, a, b) ->
                                let c2 = Cmod (a, b) in
              let s' = CondSet.remove c2 (init_set (fun x -> Cmod(a, x)) a) in
        let f x = CondSet.union (cond_and_single c1 x) in
                                let s'' = cond_and_not_single c1 (Cgeq k) in
        CondSet.fold f s' s''

(* generalization *)
let cond_and_not s1 c2 =
        let add_and_not x = CondSet.union (cond_and_not_single x c2) in
        CondSet.fold add_and_not s1 CondSet.empty

(* cond_implied_by_single c2 c1 gives true iff c1 => c2, i.e. if the set *)
(* of indexes denoted by c1 is contained in the one denoted by c2. *)
let cond_implied_by_single c2 c1 = match c1, c2 with
        | c1, c2 when c1 = c2 -> true (* shortcut *)
        | Ceq h, Ceq k -> h = k
        | Ceq h, Cgeq k 
        | Cgeq h, Cgeq k -> h >= k
        | Ceq h, Cmod (a, b) -> h mod a = b
        | Ceq h, Cgeqmod (k, a, b) -> h >= k && h mod a = b
        | Cmod (a, b), Cmod(c, d)
        | Cgeqmod(_, a, b), Cmod(c, d) -> a mod c = 0 && b mod c = d mod c
        | Cmod (a, b), Cgeqmod(k, c, d) ->
                let k' = k - k mod c + d in
                b >= k' && a mod c = 0 && b mod c = d mod c
        | Cgeqmod (h, a, b), Cgeqmod(k, c, d) ->
                let h' = h - h mod a + b in
                let k' = k - k mod c + d in
                h' >= k' && a mod c = 0 && b mod c = d mod c
        | _ -> false

(* cond_implies s1 c2 iff s1 => c2. Based on fact that (c1 || ... || ck) => c *)
(* iff c1 => c && ... && ck => c *)
let cond_implies s1 c2 = CondSet.for_all (cond_implied_by_single c2) s1

(* cond_neg_implied_by_single c2 c1 iff c1 => !c2 iff !(c1 && c2), which is *)
(* symmetric. *)
let cond_neg_implied_by_single c2 c1 = match c1, c2 with
        | Ceq h, Ceq k -> h <> k
        | Ceq h, Cgeq k | Cgeq k, Ceq h -> h < k
        | Ceq h, Cmod (a, b) | Cmod (a, b), Ceq h -> h mod a <> b
        | Ceq h, Cgeqmod (k, a, b) | Cgeqmod(k, a, b), Ceq h -> h < k || h mod a <> b
        | Cmod (a, b), Cmod (c, d)
        | Cmod (a, b), Cgeqmod(_, c, d)
        | Cgeqmod(_, a, b), Cmod (c, d)
        | Cgeqmod(_, a, b), Cgeqmod(_, c, d) ->
                let (_, _, gcd) = extended_gcd a c in
                b mod gcd <> d mod gcd
        | _ -> false

(* cond_implies_neg s1 c2 iff s1 => !c2 *)
let cond_implies_neg s1 c2 = CondSet.for_all (cond_neg_implied_by_single c2) s1

(* cond_simpl s turns Cgeqmod conditions into Cmod ones, if s implies the geq *)
(* part. *) 
let cond_simpl s = function
        | Cgeqmod(k, a, b) when cond_implies s (Cgeq k) -> Cmod(a, b)
        | c -> c 

(** Simplify the cost expression, removing useless conditions in it *) 
let remove_useless_branches =
        (* conds represents the info known while descending the branches *)
        let rec simplify' conds = function
                | Exact k -> Exact k
                | Ternary (i, gen_cond, if_true, if_false) ->
                        assert (CondSet.cardinal gen_cond = 1);
                        let cond = CondSet.choose gen_cond in
                        (* if it is the first time a condition on i is encountered, we ensure *)
                        (* that conds holds a default value that will be the "true" of these *)
                        (* conditions (i.e. { Cgeq 0 }) *)
                        ExtArray.ensure conds i;
                        let conds_i = ExtArray.get conds i in
                        (* if conds => cond, then we can erase the if_false branch *)
                        if cond_implies conds_i cond then (simplify' conds if_true) else
                        (* if conds => !cond, then we can erase if_true *)
                        if cond_implies_neg conds_i cond then (simplify' conds if_false) else
                        begin
                                let cond' = cond_simpl conds_i cond in
                                (* simplify the if_true branch knowing that cond *)
                                ExtArray.set conds i (cond_and conds_i cond);
                                let if_true' = simplify' conds if_true in
                                (* simplify the if_false branch knowing that !cond *)
                                ExtArray.set conds i (cond_and_not conds_i cond);
                                let if_false' = simplify' conds if_false in
                                Ternary (i, CondSet.singleton cond', if_true', if_false')
                        end in
        let conds = ExtArray.make ~buff:4 0 (CondSet.singleton (Cgeq 0)) in
        simplify' conds

let gen_or = CondSet.union

let gen_and s1 s2 = 
        let f c s = CondSet.union s (cond_and s1 c) in
        CondSet.fold f s2 CondSet.empty
        
let gen_and_not s1 s2 =
        let f c s = cond_and_not s c in
        CondSet.fold f s2 s1
        
let gen_not s = gen_and_not (CondSet.singleton (Cgeq 0)) s

let rec remove_useless_branchings = function
        | Exact k -> Exact k
        | Ternary (i, c1, left, right) ->
                let left = remove_useless_branchings left in
    let right = remove_useless_branchings right in
                match left, right with
                        | _, _ when left = right -> left
                        | Ternary(j, c2, lleft, lright), _
                          when i = j && lleft = right ->
                                let c = gen_or (gen_not c1) c2 in
                                Ternary(i, c, lleft, lright)
      | Ternary(j, c2, lleft, lright), _
        when i = j && lright = right ->
        let c = gen_and c1 c2 in
        Ternary(i, c, lleft, lright)
      | _, Ternary(j, c2, rleft, rright)
        when i = j && left = rleft ->
        let c = gen_or c1 c2 in
        Ternary(i, c, rleft, rright)
      | _, Ternary(j, c2, rleft, rright)
        when i = j && left = rright ->
        let c = gen_and_not c2 c1 in
        Ternary(i, c, rleft, rright)
                        | _ -> Ternary (i, c1, left, right)

let rec chose_smaller_cond = function
        | Exact k -> Exact k
        | Ternary (i, c, left, right) ->
                let left = chose_smaller_cond left in
                let right = chose_smaller_cond right in 
                let c_not = gen_not c in
                if CondSet.cardinal c > CondSet.cardinal c_not then
                        Ternary (i, c_not, right, left)
                else
                        Ternary (i, c, left, right)

let cost_expr_mapping_of_cost_mapping m =
    let sets = indexing_sets_from_cost_mapping m in
    let f at s =
                        let e = cost_mapping_ind at [] m s in
                        let e = remove_useless_branches e in
                        let e = remove_useless_branchings e in
                        let e = chose_smaller_cond e in
      Atom.Map.add at e in
    Atom.Map.fold f sets Atom.Map.empty