open Ir
open Cfg
open Dominance

module PairMap = Map.Make(struct
  type t = int * int
  let compare (a1, b1) (a2, b2) =
    match Stdlib.compare a1 a2 with
    | 0 -> Stdlib.compare b1 b2
    | c -> c
end)

type alloca_info = {
  aid : value;
  mutable stores : (value * value) list;
  mutable loads : (value * value) list;
}

let find_allocas cfg =
  let info = ref IntMap.empty in
  IntMap.iter (fun _ b ->
    List.iter (function
      | Alloca a -> info := IntMap.add a { aid = a; stores = []; loads = [] } !info
      | _ -> ()
    ) b.insts
  ) cfg.blocks;
  IntMap.iter (fun _ b ->
    List.iter (function
      | Store (v, p) ->
        (try let ai = IntMap.find p !info in ai.stores <- (v, b.id) :: ai.stores with Not_found -> ())
      | Load (r, p) ->
        (try let ai = IntMap.find p !info in ai.loads <- (r, b.id) :: ai.loads with Not_found -> ())
      | _ -> ()
    ) b.insts
  ) cfg.blocks;
  IntMap.fold (fun _ ai acc -> ai :: acc) !info []

let place_phi cfg df allocas =
  let phi_needed = ref PairMap.empty in
  List.iter (fun ai ->
    let blocks_with_store = ref (List.fold_left (fun s (_, b) -> IntSet.add b s) IntSet.empty ai.stores) in
    let work = ref (IntSet.elements !blocks_with_store) in
    while !work <> [] do
      let b = List.hd !work in
      work := List.tl !work;
      (try
        let frontiers = IntMap.find b df in
        IntSet.iter (fun f ->
          let key = (f, ai.aid) in
          if not (PairMap.mem key !phi_needed) then begin
            phi_needed := PairMap.add key true !phi_needed;
            if not (IntSet.mem f !blocks_with_store) then begin
              blocks_with_store := IntSet.add f !blocks_with_store;
              work := f :: !work
            end
          end
        ) frontiers
      with Not_found -> ())
    done
  ) allocas;
  !phi_needed

let rename cfg phi_needed allocas =
  let stacks = ref IntMap.empty in
  let push aid v =
    let s = try IntMap.find aid !stacks with Not_found -> [] in
    stacks := IntMap.add aid (v :: s) !stacks
  in
  let pop aid =
    let s = IntMap.find aid !stacks in
    if List.tl s = [] then stacks := IntMap.remove aid !stacks
    else stacks := IntMap.add aid (List.tl s) !stacks
  in
  let top aid =
    try Some (List.hd (IntMap.find aid !stacks))
    with Not_found -> None
  in
  (* For each (block, alloca) needing a phi, accumulate operands as the DFS
     visits each predecessor. Each entry maps a predecessor id to the SSA
     version of the variable on that edge. These are read only AFTER the full
     DFS completes, so every predecessor has been visited. *)
  let phi_ops = ref (PairMap.map (fun _ -> ref []) phi_needed) in
  let body_insts = ref IntMap.empty in
  let process_block id b =
    (* A block that needs a phi defines a new version of the variable (using
       the alloca id as the SSA name). Push it so subsequent loads in this
       block read the phi's result rather than a stale predecessor version. *)
    let pushed = ref [] in
    List.iter (fun ai ->
      if PairMap.mem (id, ai.aid) phi_needed then (push ai.aid ai.aid; pushed := ai.aid :: !pushed)
    ) allocas;
    let insts = ref [] in
    List.iter (fun i ->
      match i with
      | Alloca a -> ()   (* promoted away *)
      | Store (v, p) ->
        (* A store defines a new version of the variable. Push it once. *)
        push p v; pushed := p :: !pushed
      | Load (r, p) ->
        (match top p with
         | Some v -> insts := Move (r, v) :: !insts
         | None -> ())
      | Phi _ -> ()      (* input phis are replaced by the inserted ones *)
      | _ -> insts := i :: !insts
    ) b.insts;
    (* Rev to restore original order (insts was built by consing). *)
    body_insts := IntMap.add id (List.rev !insts) !body_insts;
    !pushed
  in
  let add_phi_operand block aid pred v =
    if PairMap.mem (block, aid) phi_needed then
      let ops_ref = PairMap.find (block, aid) !phi_ops in
      (* Store as (pred, value) so IntMap.of_list below builds a pred->value map. *)
      ops_ref := (pred, v) :: !ops_ref
  in
  let rec dfs id visited =
    if IntSet.mem id visited then visited
    else
      let visited = IntSet.add id visited in
      let b = IntMap.find id cfg.blocks in
      let pushed = process_block id b in
      let visited = List.fold_left (fun vis s ->
        (* Edge id -> s: the version of each variable live on entry to s is the
           current top of its stack. Record it as s's phi operand from id,
           BEFORE descending into s (so s's own phi push doesn't interfere). *)
        List.iter (fun ai ->
          match top ai.aid with
          | Some v -> add_phi_operand s ai.aid id v
          | None -> ()
        ) allocas;
        dfs s vis
      ) visited b.succs in
      (* Pop every version this block pushed before backtracking. *)
      List.iter pop pushed;
      visited
  in
  ignore (dfs cfg.entry IntSet.empty);
  (* Now build phi instructions from the fully-collected operands. *)
  let phi_insts = ref IntMap.empty in
  PairMap.iter (fun (block, aid) ops_ref ->
    let ops_map = IntMap.of_list !ops_ref in
    let b = IntMap.find block cfg.blocks in
    let phi_ops_list = List.map (fun p ->
      match IntMap.find_opt p ops_map with
      | Some v -> (v, p)
      | None -> (0, p)
    ) b.preds in
    let existing = try IntMap.find block !phi_insts with Not_found -> [] in
    phi_insts := IntMap.add block (Phi (aid, phi_ops_list) :: existing) !phi_insts
  ) !phi_ops;
  (* Final per-block instruction list: phis first, then the rewritten body. *)
  let new_insts = ref IntMap.empty in
  IntMap.iter (fun id body ->
    let phis = try IntMap.find id !phi_insts with Not_found -> [] in
    new_insts := IntMap.add id (phis @ body) !new_insts
  ) !body_insts;
  !new_insts

let mem2reg f =
  let cfg = build_cfg f in
  let doms = dominators cfg in
  let df = dominance_frontiers cfg doms in
  let allocas = find_allocas cfg in
  let phi_needed = place_phi cfg df allocas in
  let new_insts = rename cfg phi_needed allocas in
  let new_blocks = IntMap.mapi (fun id b ->
    { b with insts = try IntMap.find id new_insts with Not_found -> b.insts }
  ) cfg.blocks in
  (* Preserve the original input block order for deterministic output. *)
  let ordered = List.map (fun b -> IntMap.find b.id new_blocks) f.blocks in
  { f with blocks = ordered }