open ExtList
open ExtString

type igraph =
  {graph: Igraph.graph;
   tnode: Temp.temp -> Igraph.node;
   gtemp: Igraph.node -> Temp.temp;
   moves: (Igraph.node * Igraph.node) list}

type liveset = unit Temp.Table.t * Temp.temp list
type livemap = liveset Flow.Graph.Table.t

(* Set functions on lists *)

let rec add e = function
    [] -> [e]
  | hd :: tl ->
      if e = hd then
        hd :: tl
      else if e < hd then
        e :: hd :: tl
      else
        hd :: add e tl

let rec union x y =
  match x, y with
    [], _ -> y
  | _, [] -> x
  | hd :: tl, _ ->
      add hd (union tl y)

let rec diff x y =
  match x, y with
    [], _ -> []
  |  _, [] -> x
  | hd :: tl, _ ->
      if List.mem hd y then
        diff tl y
      else
        add hd (diff tl y)

let show_livemap out_channel livemap =
  Flow.Graph.Table.iter
    (fun fnode (_, temps) ->
       output_string out_channel (Flow.Graph.nodename fnode ^ "\t");
       output_string out_channel (String.join " " (List.map (Frame.string_of_temp) temps));
       output_string out_channel "\n")
    livemap

let liveness ({Flow.def; use; ismove; control} as flowgraph) =
  let fnodes = Flow.Graph.nodes control in
  let emptytab = List.fold_left
      (fun tab fnode -> Flow.Graph.Table.add fnode [] tab)
      Flow.Graph.Table.empty
      fnodes in
  let rec loop intab outtab =
    let (intab', outtab') = List.fold_right
        (fun fnode (intab, outtab) ->
           let use = Flow.Graph.Table.find fnode use in
           let def = Flow.Graph.Table.find fnode def in
           let liveout = Flow.Graph.Table.find fnode outtab in
           let livein' = union use (diff liveout def) in
           let liveout' = List.fold_left
               (fun liveout succ ->
                  union liveout (Flow.Graph.Table.find succ intab))
               liveout
               (Flow.Graph.succ fnode) in
           let intab' = Flow.Graph.Table.add fnode livein' intab in
           let outtab' = Flow.Graph.Table.add fnode liveout' outtab in
             (intab', outtab'))
        fnodes
        (intab, outtab) in
      if List.for_all
           (fun fnode ->
             Flow.Graph.Table.find fnode intab = Flow.Graph.Table.find fnode intab' &&
             Flow.Graph.Table.find fnode outtab = Flow.Graph.Table.find fnode outtab')
           fnodes then
        (intab', outtab')
      else
         loop intab' outtab' in
  let (_, outtab) = loop emptytab emptytab in
    List.fold_left
      (fun livemap fnode ->
         let livelist = Flow.Graph.Table.find fnode outtab in 
         let liveset = List.fold_left
             (fun tab t -> Temp.Table.add t () tab)
             Temp.Table.empty
             livelist in
           Flow.Graph.Table.add fnode (liveset, livelist) livemap)
      Flow.Graph.Table.empty
      fnodes

let interference_graph ({Flow.def; use; ismove; control} as flowgraph) =
  let livemap = liveness flowgraph in
  let igraph = Igraph.new_graph () in
  let add_inode (tnode, gtemp) t =
    if Temp.Table.mem t tnode then
      (tnode, gtemp)
    else
      let inode = Igraph.new_node igraph in
      let tnode' = Temp.Table.add t inode tnode in
      let gtemp' = Igraph.Table.add inode t gtemp in
        (tnode', gtemp') in
  let add_iedge n1 n2 =
    if not (Igraph.eq n1 n2) && not (List.exists (Igraph.eq n1) (Igraph.adj n2)) then
      Igraph.mk_edge n1 n2 in
  let rec loop tnode gtemp moves = function
      [] ->
        {graph=igraph;
         tnode=(fun t -> Temp.Table.find t tnode);
         gtemp=(fun inode -> Igraph.Table.find inode gtemp);
         moves=List.rev moves}
    | fnode :: fnodes ->
        let use = Flow.Graph.Table.find fnode use in
        let def = Flow.Graph.Table.find fnode def in
        let ismove = Flow.Graph.Table.find fnode ismove in
        let (livetab, livelist) = Flow.Graph.Table.find fnode livemap in
        let (tnode', gtemp') = List.fold_left add_inode (tnode, gtemp) (use @ def @ livelist) in
        let moves' = if ismove then
            let u = Temp.Table.find (List.hd def) tnode' in
            let v = Temp.Table.find (List.hd use) tnode' in
              (u, v) :: moves
          else
            moves in
        let adj = if ismove && Temp.Table.mem (List.hd use) livetab then
            List.filter (fun t -> t <> List.hd use) livelist
          else
            livelist in
          List.iter
            (fun d ->
               let inode = Temp.Table.find d tnode' in
               let iadj = List.map (fun t -> Temp.Table.find t tnode') adj in
                 List.iter (add_iedge inode) iadj)
            def;
          loop tnode' gtemp' moves' fnodes in
  let fnodes = Flow.Graph.nodes control in
  let igraph = loop Temp.Table.empty Igraph.Table.empty [] fnodes in
  let liveout fnode = snd (Flow.Graph.Table.find fnode livemap) in
    (igraph, liveout)

let show out_channel {graph; gtemp; _} =
  let say = output_string out_channel in
    say "\n";
  List.iter
    (fun n ->
      say ((Frame.string_of_temp (gtemp n)) ^ ":\t" ^
           (String.join ", " (List.map (fun gn -> Frame.string_of_temp (gtemp gn)) (Igraph.adj n))) ^
           ": " ^
           (string_of_int (List.length (Igraph.adj n))) ^
           "\n"))
    (Igraph.nodes graph)