open Ir

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

type cfg = {
  blocks : block IntMap.t;
  entry : value;
  mutable defs : inst IntMap.t;
}

let build_cfg (f : func) =
  let blocks = List.fold_left (fun m b -> IntMap.add b.id b m) IntMap.empty f.blocks in
  let cfg = { blocks; entry = f.entry; defs = IntMap.empty } in
  let add_edges b =
    let rec scan acc = function
      | [] -> List.rev acc
      | Ret _ :: _ -> List.rev acc
      | BinOp (_, _, t) :: rest -> scan (t :: acc) rest
      | Br t :: rest -> scan (t :: acc) rest
      | CondBr (_, t1, t2) :: rest -> scan (t2 :: t1 :: acc) rest
      | _ :: rest -> scan acc rest
    in
    b.succs <- scan [] b.insts
  in
  IntMap.iter (fun _ b -> add_edges b) blocks;
  IntMap.iter (fun _ b ->
    List.iter (fun s ->
      try
        let sb = IntMap.find s cfg.blocks in
        sb.preds <- b.id :: sb.preds
      with Not_found -> ()
    ) b.succs
  ) blocks;
  let defs = IntMap.fold (fun _ b acc ->
    List.fold_left (fun d inst ->
      match inst with
      | Alloca a -> IntMap.add a inst acc
      | Load (r, _) -> IntMap.add r inst d
      | BinOp (r, _, _) -> IntMap.add r inst d
      | Move (r, _) -> IntMap.add r inst d
      | Phi (r, _) -> IntMap.add r inst d
      | _ -> d
    ) acc b.insts
  ) blocks IntMap.empty in
  cfg.defs <- defs;
  cfg

let get_block cfg id = IntMap.find id cfg.blocks
let get_succs cfg id = (get_block cfg id).succs
let get_preds cfg id = (get_block cfg id).preds
let get_def cfg v = IntMap.find v cfg.defs