open Ir
open Cfg

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

(* Immediate dominator, derived correctly from full dominator sets.
   The idom of n is the unique proper dominator d of n that is itself
   dominated by every other proper dominator of n (i.e. the deepest one).
   This replaces the previous ad-hoc heuristic, which could mis-pick an outer
   dominator when a node had 3+ dominators. *)
let idom cfg 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
    (* Find the d in strict_doms that every other strict dominator dominates. *)
    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 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 all = IntMap.fold (fun k _ s -> IntSet.add k s) cfg.blocks IntSet.empty in
  let init = IntSet.singleton cfg.entry in
  let rpo = rev_postorder cfg in
  let rec fixpoint doms =
    let new_doms = List.fold_left (fun acc id ->
      if id = cfg.entry then IntMap.add id init acc
      else
        let b = IntMap.find id cfg.blocks in
        match b.preds with
        | [] -> IntMap.add id IntSet.empty acc
        | p :: ps ->
          let inter = List.fold_left (fun a pr ->
            try IntSet.inter a (IntMap.find pr acc)
            with Not_found -> a
          ) (try IntMap.find p acc with Not_found -> IntSet.empty) ps in
          IntMap.add id (IntSet.add id inter) acc
    ) IntMap.empty rpo in
    if IntMap.equal IntSet.equal doms new_doms then doms
    else fixpoint new_doms
  in
  fixpoint (IntMap.map (fun _ -> all) cfg.blocks)

let dominates doms a b =
  IntSet.mem a (IntMap.find b 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 cfg doms id in
      List.iter (fun p ->
        (* Runner p is a predecessor of Y (p <> Y). Add Y to DF at each node
           from p up to (but not including) Y's idom. p itself is included. *)
        let rec walk cur =
          if Some cur = idom_y then ()   (* reached idom(Y): stop, don't add *)
          else begin
            df := IntMap.add cur (IntSet.add id (IntMap.find cur !df)) !df;
            (match idom cfg doms cur with
             | None -> ()
             | Some n -> walk n)
          end
        in
        walk p
      ) b.preds
  ) cfg.blocks;
  !df