open Ir
open Cfg

module IntMap = Map.Make(Int)
module IntSet = Set.Make(Int)

let rev_postorder cfg =
  let visited = ref IntSet.empty in
  let order = ref [] in
  let rec dfs id =
    if IntSet.mem id !visited then ()
    else begin
      visited := IntSet.add id !visited;
      let b = IntMap.find id cfg.blocks in
      List.iter dfs b.succs;
      order := id :: !order
    end
  in
  dfs cfg.entry;
  let remaining = IntMap.fold (fun k _ acc ->
    if IntSet.mem k !visited then acc else k :: acc
  ) cfg.blocks [] in
  !order @ remaining

let dominators cfg =
  let rpo = rev_postorder cfg in
  let rpo_idx = Hashtbl.create 64 in
  List.iteri (fun i x -> Hashtbl.add rpo_idx x i) rpo;
  let idom = ref IntMap.empty in
  let get_rank v =
    try Hashtbl.find rpo_idx v with Not_found -> max_int in
  List.iter (fun w ->
    if w = cfg.entry then ()
    else
      let b = IntMap.find w cfg.blocks in
      match b.preds with
      | [] -> ()
      | p0 :: ps ->
        let find_best v =
          try IntMap.find v !idom with Not_found -> v
        in
        let best = ref (find_best p0) in
        let best_rank = ref (get_rank !best) in
        List.iter (fun p ->
          let candidate = find_best p in
          let r = get_rank candidate in
          if r < !best_rank then begin best := candidate; best_rank := r end
        ) ps;
        let semi_w = !best in
        let rec walk cur best_id =
          if cur = w then best_id
          else
            let r = get_rank cur in
            let best_id = if r < get_rank best_id then cur else best_id in
            try
              let parent = IntMap.find cur !idom in
              walk parent best_id
            with Not_found -> best_id
        in
        let idom_w = walk semi_w w in
        idom := IntMap.add w idom_w !idom;
        ignore semi_w
  ) rpo;
  let doms = ref IntMap.empty in
  List.iter (fun id ->
    let rec collect v acc =
      if IntSet.mem v acc then acc
      else
        let acc = IntSet.add v acc in
        try collect (IntMap.find v !idom) acc with Not_found -> acc
    in
    doms := IntMap.add id (collect id IntSet.empty) !doms
  ) rpo;
  !doms

let dominates doms a b =
  IntSet.mem a (IntMap.find b doms)

let idom_from_doms doms id =
  let s = IntMap.find id doms in
  let without = IntSet.diff s (IntSet.singleton id) in
  if IntSet.is_empty without then None
  else
    let strict_doms = IntSet.elements without in
    let rec find = function
      | [] -> None
      | d :: rest ->
        if List.for_all (fun m -> m = d || IntSet.mem d (IntMap.find m doms)) strict_doms
        then Some d
        else find rest
    in
    find strict_doms

let dominance_frontiers cfg doms =
  let df = ref IntMap.empty in
  IntMap.iter (fun id _ -> df := IntMap.add id IntSet.empty !df) cfg.blocks;
  IntMap.iter (fun id b ->
    match b.preds with
    | [] | [_] -> ()
    | _ ->
      let idom_y = idom_from_doms doms id in
      List.iter (fun p ->
        let rec walk cur =
          if Some cur = idom_y then ()
          else begin
            df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
            (match idom_from_doms doms cur with
             | None -> ()
             | Some n -> walk n)
          end
        in
        walk p
      ) b.preds
  ) cfg.blocks;
  !df