open Xi_lib
open Ir
open Ir_utils

module Make() = struct

  module Implementation(M:sig
                            val cfg: ControlFlowGraph.t
                            val initial: Analysis_domain.ConstantFolding.domain
                          end) = struct

    open M
    open Analysis
    let initialize_table () =
      let table = Hashtbl.create 513 in
      let kw = Knowledge.make ~pre:initial ~post:initial in
      let blk_kw =  BlockKnowledge.make_simple kw in
      let set v = Hashtbl.replace table v blk_kw in
      List.iter set @@ ControlFlowGraph.labels cfg;
      table
    let result = initialize_table ()

    let join_elem a b =match (a,b) with
      | None,None -> None
      | Some a,Some b -> if a=b then Some a else Some None
      | None, Some _ -> Some None
      | Some _, None -> Some None

    let int32_of_bool b = if b then Int32.one else Int32.zero
     
    let transfer_instr instr input =
      let get v = match v with
        | E_Int x -> Some x
        | E_Reg r -> (match Ir.RegMap.find_opt r input with
                      |None -> None
                      |Some None -> None
                      |Some Some x ->Some x)
      in
      let aux r l1 l2 op =
        (let v= (match get l1,get l2
                 with
                 |Some a,Some b -> op a b
                 |_ ->None
                )
         in match v
            with
            | None ->input
            | Some v -> Ir.RegMap.add r (Some v) input)
        
      in
      let output = match instr with
        | I_Add (r,l1,l2) -> aux r l1 l2 (fun a b -> Some (Int32.add a b))
        | I_Sub (r,l1,l2) -> aux r l1 l2 (fun a b -> Some (Int32.sub a b))
        | I_Div (r,l1,l2) -> aux r l1 l2 (fun a b -> if b=Int32.zero then None else Some(Int32.div a b))
        | I_Mul (r,l1,l2) -> aux r l1 l2 (fun a b -> Some (Int32.mul a b))
        
        | I_Neg (r,l1) -> aux r l1 l1 (fun a b -> Some (Int32.neg a))
        | I_Not (r,l1) -> aux r l1 l1 (fun a b -> Some (Int32.lognot a))
                        
        | I_And (r,l1,l2) -> aux r l1 l2 (fun a b -> Some (Int32.logand a b))
        | I_Or (r,l1,l2) -> aux r l1 l2 (fun a b -> Some (Int32.logor a b))
        | I_Xor (r,l1,l2) -> aux r l1 l2 (fun a b -> Some (Int32.logxor a b))
                           
        | I_Set (r,cond,l1,l2) -> aux r l1 l2
                                    (fun a b -> let cmp=Int32.compare a b in
                                                Some (int32_of_bool(match cond
                                                  with
                                                  |COND_Eq -> cmp=0
                                                  |COND_Ne -> not (cmp=0)
                                                  |COND_Lt -> cmp<0
                                                  |COND_Gt -> cmp>0
                                                  |COND_Le -> cmp<=0
                                                  |COND_Ge -> cmp>=0
                                                         )))                   
        | _ -> List.fold_left (fun set elem -> Ir.RegMap.add elem None set) input @@ Ir_utils.used_registers_instr instr
             
      in
      (Knowledge.make ~pre:input ~post:output),instr

    let transfer_terminator t input =
      (Knowledge.make ~pre:input ~post:input),t

    let rec transfer_instr_list instr input = match instr with
      | [] -> [],input
      | x::xs -> 
         let y= transfer_instr x input in
         let input = Knowledge.post @@ fst y in
         let ys,output = transfer_instr_list xs input in
         (y::ys,output)
         
    let transfer_basic_block l current_knowledge =
      if l = ControlFlowGraph.exit_label cfg then current_knowledge else
        if l = ControlFlowGraph.entry_label cfg then current_knowledge else
          let t = ControlFlowGraph.terminator cfg l in
          let instr = ControlFlowGraph.block cfg l in
          let input = BlockKnowledge.pre current_knowledge in
          let instr_res,input2 = transfer_instr_list instr input in
          let t_res = transfer_terminator t input2 in
          let output = Knowledge.post @@ fst t_res in
          let block = Knowledge.make ~pre:input ~post:output in
          BlockKnowledge.make_complex ~block ~body:instr_res ~terminator:t_res
    let changed = ref true

    let join a b = Ir.RegMap.merge (fun _ x y -> join_elem x y ) a b

    let update_in label =
      let input = List.fold_left join initial (
                      (List.map (fun l -> BlockKnowledge.post (Hashtbl.find result l)))
                        (ControlFlowGraph.predecessors cfg label))
      in
      let old = (Hashtbl.find result label) in
      let old_input = BlockKnowledge.pre old in
      if Ir.RegMap.compare compare input old_input = 0 then () else
        changed:=true;
      let new_kn = BlockKnowledge.alter_prepost ~pre:input old in
      Hashtbl.replace result label new_kn
      
    let update_ins () =
      let labels = ControlFlowGraph.labels cfg in
      List.iter update_in labels

      
    let update_out label =
      let old = (Hashtbl.find result label) in
      let old_input = BlockKnowledge.post old in
      let new_kn = transfer_basic_block label old in
      if Ir.RegMap.compare compare (BlockKnowledge.post new_kn) old_input = 0 then () else
        changed:=true;
      Hashtbl.replace result label new_kn
      
    let update_outs () =
      let labels = ControlFlowGraph.labels cfg in
      List.iter update_out labels
      
    let rec compute_fixpoint () =
      changed := false;
      update_ins ();
      update_outs ();
      if not !changed then () else
        compute_fixpoint ()
      
      

    let analyse () =
      compute_fixpoint ();
      result

      
  end



  (* Skontruuj wartość ekstremalną *)
  let rec make_initial n =
    if n = 0 then Ir.RegMap.empty else
      Ir.RegMap.add (Ir.REG_Tmp (n-1)) None @@ make_initial (n-1)
    

  let analyse proc : Xi_lib.Analysis_domain.ConstantFolding.table =
    let initial = make_initial @@ Ir.formal_parameters_of_procedure proc in
    let cfg = Ir.cfg_of_procedure proc in
    let module Instance = Implementation(struct let cfg = cfg let initial = initial end) in 
    let result = Instance.analyse () in
    result

end