module type PrintableType = sig
  type t
  val pp : out_channel -> t -> unit
  val pp_list : out_channel -> t list -> unit
end

let globalIdNode = ref 0;

module Node = struct
  type t = {
    id: int;
  }
  let compare a b = compare a.id b.id

  let create () =
    globalIdNode := !globalIdNode + 1;
    {id = !globalIdNode;}
end

module NodeMap = struct
  include Map.Make(Node)

  (* adds the input to the tail of the list for the associated node *)
  let add_to_list_last x data m =
    let add = function None -> Some [data]
                     | Some l -> Some (l @ [data]) in
    update x add m
end

module NodeSet = Set.Make(Node)

type 'a cfginternal = {
  empty: bool;
  nodes: NodeSet.t;
  edges: (Node.t * (Node.t option)) NodeMap.t;
  reverse_edges: (Node.t list) NodeMap.t;
  input_val: int option;
  input_output_var: (string * string) option;
  initial: Node.t option;
  terminal: Node.t option;
  content: 'a list NodeMap.t;
}

module type C = sig
  type elt
  type t = elt cfginternal

  val empty : t
  val merge : t -> t -> Node.t -> Node.t -> t
  val concat : t -> t -> t
  val add_to_last_node : elt -> t -> t

  val pp : out_channel -> t -> unit
end

module Make (M: PrintableType) = struct
  type elt = M.t
  type t = elt cfginternal

  let empty : t =
    { empty = true;
      nodes = NodeSet.empty;
      edges = NodeMap.empty;
      reverse_edges = NodeMap.empty;
      input_val = None;
      input_output_var = None;
      initial = None;
      terminal = None;
      content = NodeMap.empty }

  let merge (cfg1: t) (cfg2: t) (entry_node: Node.t) (exit_node: Node.t) : t =
    match (cfg1.empty, cfg2.empty) with
      true, _ -> cfg2
    | _, true -> cfg1
    | false, false ->
      let cfg1_initial = Option.get cfg1.initial in
      let cfg2_initial = Option.get cfg2.initial in
      let cfg1_terminal = Option.get cfg1.terminal in
      let cfg2_terminal = Option.get cfg2.terminal in
      { empty = false;
        nodes = NodeSet.union cfg1.nodes cfg2.nodes |>
                NodeSet.add entry_node |>
                NodeSet.add exit_node;
        edges = NodeMap.union
            (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.edges cfg2.edges |>
                NodeMap.add entry_node (cfg1_initial, Some cfg2_initial) |>
                NodeMap.add cfg1_terminal (exit_node, None) |>
                NodeMap.add cfg2_terminal (exit_node, None);
        reverse_edges = NodeMap.union
            (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.reverse_edges cfg2.reverse_edges |>
                       NodeMap.add_to_list cfg1_initial entry_node |>
                       NodeMap.add_to_list cfg2_initial entry_node |>
                       NodeMap.add_to_list exit_node cfg1_terminal |>
                       NodeMap.add_to_list exit_node cfg2_terminal;
        input_val = cfg1.input_val;
        input_output_var = cfg1.input_output_var;
        initial = Some entry_node;
        terminal = Some exit_node;
        content = NodeMap.union
            (fun _ -> failwith "Failed merging code of cfg.")
            cfg1.content cfg2.content
      }

  let concat (cfg1: t) (cfg2: t) : t =
    match (cfg1.empty, cfg2.empty) with
      true, _ -> cfg2
    | _, true -> cfg1
    | false, false ->
      let cfg1_initial = Option.get cfg1.initial in
      let cfg2_initial = Option.get cfg2.initial in
      let cfg1_terminal = Option.get cfg1.terminal in
      let cfg2_terminal = Option.get cfg2.terminal in
      { empty = false;
        nodes = NodeSet.union cfg1.nodes cfg2.nodes;
        edges = NodeMap.union
            (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.edges cfg2.edges |>
                NodeMap.add cfg1_terminal (cfg2_initial, None);
        reverse_edges = NodeMap.union
            (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.reverse_edges cfg2.reverse_edges |>
                       NodeMap.add_to_list cfg2_initial cfg1_terminal;
        input_val = cfg1.input_val;
        input_output_var = cfg1.input_output_var;
        initial = Some cfg1_initial;
        terminal = Some cfg2_terminal;
        content = NodeMap.union
            (fun _ -> failwith "Failed merging code of cfg.")
            cfg1.content cfg2.content
      }

  let add_to_last_node (new_content: elt) (cfg: t) : t =
    if cfg.empty then
      let new_node = Node.create () in
      { empty = false;
        nodes = NodeSet.singleton new_node;
        edges = NodeMap.empty;
        reverse_edges = NodeMap.empty;
        input_val = None;
        input_output_var = None;
        initial = Some new_node;
        terminal = Some new_node;
        content = NodeMap.singleton new_node [new_content]
      }
    else
      let prevcfg_terminal = Option.get cfg.terminal in
      { cfg with
        content = (NodeMap.add_to_list_last
                  prevcfg_terminal
                  new_content
                  cfg.content) }

  let pp (ppf) (c: t) : unit =
    Printf.fprintf ppf "Nodes' ids: ";
    List.iter
      (fun (x : Node.t) -> Printf.fprintf ppf "%d " x.id)
      (NodeSet.to_list c.nodes);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Nodes' edges:\n";
    List.iter
      (fun ((n, (a, b)) : (Node.t * (Node.t * Node.t option))) : unit ->
         match b with
           None -> Printf.fprintf ppf "\t%d -> %d\n" n.id a.id
         | Some b -> Printf.fprintf ppf "\t%d -> %d, %d\n" n.id a.id b.id )
      (NodeMap.to_list c.edges);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Nodes' back edges:\n";
    List.iter (fun ((n, xs) : (Node.t * (Node.t list))) : unit ->
        Printf.fprintf ppf "\t%d -> " n.id;
        List.iter (fun (x: Node.t) -> Printf.fprintf ppf "%d, " x.id) xs;
        Printf.fprintf ppf "\n"
      ) (NodeMap.to_list c.reverse_edges);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Input Value: ";
    (match c.input_val with
       Some i -> Printf.fprintf ppf "%d" i;
     | None   -> Printf.fprintf ppf "None";);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Input and Output Vars: ";
    (match c.input_output_var with
       Some (i, o) -> Printf.fprintf ppf "(in: %s, out: %s)" i o;
     | None   -> Printf.fprintf ppf "None";);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Initial node's id: ";
    (match c.initial with
       Some i -> Printf.fprintf ppf "%d" (i.id);
     | None   -> Printf.fprintf ppf "None";);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Terminal node's id: ";
    (match c.terminal with
       Some i -> Printf.fprintf ppf "%d" (i.id);
     | None   -> Printf.fprintf ppf "None";);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Code:\n";
    List.iter (fun ((n, stms) : Node.t * elt list) : unit ->
        Printf.fprintf ppf "\tid %d --> %a\n%!" n.id M.pp_list stms
      ) (NodeMap.to_list c.content);
    Printf.fprintf ppf "\n";
end
;;