open Xi_lib
open Iface
open Ir
open Ir_utils
open Analysis (* <--- tu mogą być pomocne komentarze *)
open Analysis_domain

module Make() = struct

    module Implementation(M:sig val cfg: ControlFlowGraph.t end) = struct
        open M

        (*
         * Zwróćmy tablicę gdzie każdy wierzchołek jest zainicjalizowany na
         * konstruktor Simple (typ BlockKnowledge) gdzie na wejściu/wyjściu
         * bloku mamy pusty zbiór rejestrów.
         *
         * Wierzchołki oznaczające basic-bloki powinny ostatecznie być opisane
         * konstruktorem Complex, ale początkowo dla wygody możemy ustawić je na Simple.
         * Ważne aby funkcja przeliczająca wiedzę dla bloku podstawowego ostatecznie
         * opisał blok za pomocą konstruktora Complex. 
         *)
        let initialize_table () =
            let table = Hashtbl.create 513 in
            let kw = Knowledge.make ~pre:RegSet.empty ~post:RegSet.empty 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 union set list=
          List.fold_left (fun x y->RegSet.add y x) set list
          
        let diff set list =
          List.fold_left (fun x y->RegSet.remove y x) set list
          
        let result : LiveVariables.table = initialize_table ()

        let transfer_instr instr input =
          let output = union (diff input (Ir_utils.defined_registers_instr instr)) (Ir_utils.used_registers_instr instr) in
          (Knowledge.make ~pre:output ~post:input),instr

          let transfer_terminator t input =
          let output = union (diff input (Ir_utils.defined_registers_terminator t)) (Ir_utils.used_registers_terminator t) in
          (Knowledge.make ~pre:output ~post:input),t

          let rec transfer_instr_list instr input = match instr with
            | [] -> [],input
            | x::xs -> let ys,input = transfer_instr_list xs input in
                       let y= transfer_instr x input in
                       let output = Knowledge.pre (fst y) 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.post current_knowledge in
          let t_res = transfer_terminator t input in
          let input2 = Knowledge.pre @@ fst t_res in
          let instr_res,output = transfer_instr_list instr input2 in
          let block = Knowledge.make ~pre:output ~post:input in
          BlockKnowledge.make_complex ~block ~body:instr_res ~terminator:t_res
        let changed = ref true

        let update_in label =
          let input = List.fold_left RegSet.union RegSet.empty
                        (List.map (fun l -> BlockKnowledge.pre (Hashtbl.find result l))
                           (ControlFlowGraph.successors cfg label) ) in
          let old = (Hashtbl.find result label) in
          let old_input = BlockKnowledge.post old in
          if input=old_input then () else
            changed:=true;
          let new_kn = BlockKnowledge.alter_prepost ~post: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 new_kn=old 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

    let analyse cfg = 
        let module Instance = Implementation(struct let cfg = cfg end) in
        Instance.analyse ()


end