module type PrintableType = sig
  type t
  val pp : out_channel -> t -> unit
  val pplist : 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)

  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;
  reverseEdges: (Node.t list) NodeMap.t;
  inputVal: int option;
  inputOutputVar: (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 addToLastNode : 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;
      reverseEdges = NodeMap.empty;
      inputVal = None;
      inputOutputVar = None;
      initial = None;
      terminal = None;
      content = NodeMap.empty }

  let merge (cfg1: t) (cfg2: t) (entryNode: Node.t) (exitNode: Node.t) : t =
    match (cfg1.empty, cfg2.empty) with
      true, _ -> cfg2
    | _, true -> cfg1
    | false, false ->
      let cfg1initial = Option.get cfg1.initial in
      let cfg2initial = Option.get cfg2.initial in
      let cfg1terminal = Option.get cfg1.terminal in
      let cfg2terminal = Option.get cfg2.terminal in
      { empty = false;
        nodes = NodeSet.union cfg1.nodes cfg2.nodes |>
                NodeSet.add entryNode |>
                NodeSet.add exitNode;
        edges = NodeMap.union (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.edges cfg2.edges |>
                NodeMap.add entryNode (cfg1initial, Some cfg2initial) |>
                NodeMap.add cfg1terminal (exitNode, None) |>
                NodeMap.add cfg2terminal (exitNode, None);
        reverseEdges = NodeMap.union (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.reverseEdges cfg2.reverseEdges |>
                       NodeMap.add_to_list cfg1initial entryNode |>
                       NodeMap.add_to_list cfg2initial entryNode |>
                       NodeMap.add_to_list exitNode cfg1terminal |>
                       NodeMap.add_to_list exitNode cfg2terminal;
        inputVal = cfg1.inputVal;
        inputOutputVar = cfg1.inputOutputVar;
        initial = Some entryNode;
        terminal = Some exitNode;
        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 cfg1initial = Option.get cfg1.initial in
      let cfg2initial = Option.get cfg2.initial in
      let cfg1terminal = Option.get cfg1.terminal in
      let cfg2terminal = 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 cfg1terminal (cfg2initial, None);
        reverseEdges = NodeMap.union (fun _ -> failwith "Failed merging edges of cfg.")
            cfg1.reverseEdges cfg2.reverseEdges |>
                       NodeMap.add_to_list cfg2initial cfg1terminal;
        inputVal = cfg1.inputVal;
        inputOutputVar = cfg1.inputOutputVar;
        initial = Some cfg1initial;
        terminal = Some cfg2terminal;
        content = NodeMap.union (fun _ -> failwith "Failed merging code of cfg.")
            cfg1.content cfg2.content
      }

  let addToLastNode (newcontent: elt) (cfg: t) : t =
    match cfg.empty with
    | true -> let newnode = Node.create () in
      { empty = false;
        nodes = NodeSet.singleton newnode;
        edges = NodeMap.empty;
        reverseEdges = NodeMap.empty;
        inputVal = None;
        inputOutputVar = None;
        initial = Some newnode;
        terminal = Some newnode;
        content = NodeMap.singleton newnode [newcontent]
      }
    | false ->
      let prevcfgterminal = Option.get cfg.terminal in
      { cfg with
        content = (NodeMap.add_to_list_last
                  prevcfgterminal
                  newcontent
                  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.reverseEdges);
    Printf.fprintf ppf "\n";

    Printf.fprintf ppf "Input Value: ";
    (match c.inputVal 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.inputOutputVar 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.pplist stms
      ) (NodeMap.to_list c.content);
    Printf.fprintf ppf "\n";
end
;;