open Syntax
open Opt_tmp
;;

exception No_one_state_case_rule
exception No_progress_done

let rec remove_doublons = function
    [] -> []
  | [e] -> [e]
  | a::b::rest -> (
      if a=b then remove_doublons (b::rest)
      else (a::(remove_doublons (b::rest))))

let sort_and_remove_doublons l =
  remove_doublons (Sort.list (<) l )

let rec optimize = function
    [] -> []
  | Rule (sl, st)::r -> Rule (sl, optimize_statement st) :: (optimize r) 

and optimize_statement = function
    If(c, st, el, st') -> If(c, optimize_statement st, List.map optimize_statement_elsif el, optimize_statement st')
  | Decision(ns, d) as a -> a
  | Case(s, al, st) -> optimize_case_same_statement (Case (s, List.map optimize_statement_arm al, optimize_statement st))
    
and optimize_statement_elsif = function
    Elseif(c, st) -> Elseif(c, optimize_statement st)

and optimize_statement_arm = function
    Arm(il, st) -> Arm(il, optimize_statement st)

and optimize_case_same_statement = function
    Case(s, al, dst) -> (
      let taille = (List.length al) * 3 + 2 in
      let h = Hashtbl.create taille in
      let f x = 
	( match x with Arm(il, ast) -> 
	  try 
	    let tmp = Hashtbl.find h ast in
	    Hashtbl.remove h ast;
	    Hashtbl.add h ast (il::tmp)
	  with Not_found ->
	    Hashtbl.add h ast (il::[])
	)
      in
      List.iter f al;  
      (* we remove the default entry which is factored by the default case *)
      Hashtbl.remove h dst;
      let new_al = ref [] in
      Hashtbl.iter 
	(fun st' y -> new_al := (Arm ((sort_and_remove_doublons (List.flatten y)),st')::!new_al))
	h;
      Case(s, !new_al, dst) 
     )
	
let rec expand_an_arm st = function
    [] -> []
  | i::r -> (Arm([i], st) :: expand_an_arm st r)

let rec expand_arm = function
    [] -> []
  | Arm(il,st)::rest -> ((expand_an_arm st il)@(expand_arm rest))

let expand_and_sort_arm al =
  let al' = expand_arm al in
  Sort.list (fun x y -> (match (x,y) with (Arm([a],st1),Arm([b],st2)) -> (a<b))) al'

let rec create_equals_state_list = function
    [] -> []
  | i::r -> (Equals("state", i)::create_equals_state_list r)

let create_arm a st1 st2 state1 state2 =
  if (st1 = st2) then Arm([a], st1)
  else Arm([a], If(Or(create_equals_state_list state1), st1, [], st2))

let rec merge_arm al al' dst1 dst2 state1 state2 =
  ( match (al,al') with
  | ([],[]) -> []
  | (Arm([a],st)::r,[]) -> (create_arm a st dst2 state1 state2 :: (merge_arm r [] dst1 dst2 state1 state2))
  | ([], Arm([b],st)::r) -> (create_arm b dst1 st state1 state2 :: (merge_arm [] r dst1 dst2 state1 state2))
  | ((Arm([a],st)::r as r1), (Arm([b],st')::r' as r2)) -> (
      if a=b then (create_arm a st st' state1 state2 :: (merge_arm r r' dst1 dst2 state1 state2))
      else if a<b then (create_arm a st dst2 state1 state2 :: (merge_arm r r2 dst1 dst2 state1 state2))
      else (create_arm b dst1 st' state1 state2 :: (merge_arm r1 r' dst1 dst2 state1 state2))
     )
  )

(* gather_cases returns a case *)
(* state1 and state2 are state lists *)
let gather_cases varname al1 al2 dst1 dst2 state1 state2=
  let al'1 = expand_and_sort_arm al1 in
  let al'2 = expand_and_sort_arm al2 in
  let new_al = merge_arm al'1 al'2 dst1 dst2 state1 state2 in
  let new_dst = (
    if (dst1 = dst2) then dst1
    else If(Or(create_equals_state_list state1), dst1, [], dst2))
  in
  optimize_case_same_statement (Case(varname, new_al, new_dst))

(* gather_cases_from_rules must be called with rules that are cases on same var *)
let gather_cases_from_rules r1 r2 =
  match (r1,r2) with
  | (Rule(i1,Case(s1,al1,dst1)),Rule(i2,Case(s2,al2,dst2))) -> Rule (i1@i2, gather_cases s1 al1 al2 dst1 dst2 i1 i2)
	
let rec really_find_n_state_case n bl = function
  | [] -> raise No_one_state_case_rule
  | (Rule(il, Case(s,al,def)) as a)::r -> (
      if (List.length il) = n then (bl, r, a, s, al, def, il)
      else really_find_n_state_case n (a::bl) r)
  | h :: t -> really_find_n_state_case n (h::bl) t

let find_case_with_n_state n character =
  really_find_n_state_case n [] character

let main_controle_gather n character =
let (list_begin, list_end, case_as_rule_to_study, var, arm_list, default, state_list) = find_case_with_n_state n character in
let base_cout = Cout.cout_mem_rule case_as_rule_to_study in
let best_case = ref case_as_rule_to_study (* will not be used before a new affectation *) in
let best_merge = ref case_as_rule_to_study (* will not be used before a new affectation *) in
let gain_max = ref 0 in
let rec controle_gather fait = function
    [] -> 
      if (!gain_max > 0) (* we did better *)
      then (!best_merge :: fait) (* case and best_case are not in fait *)
      else raise No_progress_done 
  | (Rule(sl, Case(s,al,st)) as h)::t ->
      if (var = s) then (
	let tmp_c = gather_cases_from_rules case_as_rule_to_study h in
	let gain = base_cout + Cout.cout_mem_rule h - Cout.cout_mem_rule tmp_c in
	let old_gain = !gain_max in
	if (gain > old_gain) then (
	  let old_best = !best_case in
	  best_case := h;
	  best_merge := tmp_c;
	  gain_max := gain;
	  if (old_gain > 0) (* we already have an old_best *)
	  then controle_gather (old_best::fait) t
	  else controle_gather fait t (* called only first time *)
	 )
	else controle_gather (h::fait) t
       )
      else controle_gather (h::fait) t
  | h::t -> controle_gather (h::fait) t
in
controle_gather [] (list_begin@list_end)

let rec apply_map_on_rules f = function
    [] -> []
  | (Rule (sl, st))::r -> (Rule (sl, map_statement f st) :: apply_map_on_rules f r)

let optimize_same_if = function
  | (If(c, st, [], If(c', st', el, st'')) as a)-> 
      if st = st' then If(Or(c::[c']), st, el, st'')
      else a
  | a -> a