source: Deliverables/D2.2/8051/src/utilities/graphUtilities.ml

module type GraphType = sig

  type node
  type statement

  module NodeMap : Map.S with type key = node
  module NodeSet : Set.S with type elt = node

  type t = statement NodeMap.t

  val successors : statement -> node list
  val skip : node -> statement
  val fill_succs : statement -> node list -> statement

end

module Util (G : GraphType) = struct

  open G

  let compute_predecessor_lists (graph : G.t) : G.node list G.NodeMap.t =
    let add_to_preds pred map lbl =
      let preds =
        try
          pred :: NodeMap.find lbl map
        with
          | Not_found -> [pred] in
      NodeMap.add lbl preds map in
    let add_predecessors lbl stmt map =
    (* make sure the domain of the map will be equal to dom graph, adding *)
    (* empty sets if need be *)
      let map = if NodeMap.mem lbl map then map else
          NodeMap.add lbl [] map in
      List.fold_left (add_to_preds lbl) map (successors stmt) in
    NodeMap.fold add_predecessors graph NodeMap.empty

  let dfs_fold
      (f : node -> statement -> 'a -> 'a)
      (g : t)
      (entry : node)
      (init : 'a)
      : 'a =
  if not (NodeMap.mem entry g) then
    invalid_arg "dfs_fold: entry is not in graph"
  else
    let rec process done_set = function
      | [] -> init
      | next :: worklist when NodeSet.mem next done_set ->
        process done_set worklist
      | next :: worklist ->
        let stmt = NodeMap.find next g in
        let succs = successors stmt in
        f next stmt (process (NodeSet.add next done_set) (succs @ worklist)) in
    process NodeSet.empty [entry]

  let dfs_iter
    (f : node -> statement -> unit)
    (g : t)
    (entry : node)
    : unit =
  if not (NodeMap.mem entry g) then
    invalid_arg "dfs_iter: entry is not in graph"
  else
  let rec process done_set = function
    | [] -> ();
    | next :: worklist when NodeSet.mem next done_set ->
      process done_set worklist
    | next :: worklist ->
      let stmt = NodeMap.find next g in
      let succs = successors stmt in
      f next stmt;
      process (NodeSet.add next done_set) (succs @ worklist) in
  process NodeSet.empty [entry]

  let dead_code_elim
      (g : t)
      (entry : node)
      : t =
    let add lbl _ = NodeSet.add lbl in
    let reachable = dfs_fold add g entry NodeSet.empty in
    let is_reachable x _ = NodeSet.mem x reachable in
    NodeMap.filter is_reachable g

  let rec put_rev freshl stmts src dests graph =
    match stmts, dests with
      | [], [next] -> (src, NodeMap.add src (skip next) graph)
      | [last], _ ->
        (src, NodeMap.add src (fill_succs last dests) graph)
      | last :: stmts, _ ->
        let new_l = freshl () in
        let graph = NodeMap.add new_l (fill_succs last dests) graph in
        (new_l, snd (put_rev freshl stmts src [new_l] graph))
      | _ ->
        invalid_arg "successors of statement and translation do not match"

  let replace freshl lbl stmts succs g =
    let succs = match succs with
      | Some x -> x
      | None -> successors (NodeMap.find lbl g) in

    snd (put_rev freshl (List.rev stmts) lbl succs g)

  let replace'
      (freshl : unit -> node)
      (freshr : 'u -> 't -> 'u * 'r)
      (lbl : node)
      (stmts : (statement, 'r, 't) BList.t)
      (succs : node list option)
      (def : 'u)
      (g : t)
      : 'u * t =
    let succs = match succs with
      | Some x -> x
      | None -> successors (NodeMap.find lbl g) in
    let (def, l) = BList.compile freshr def stmts in
    (def, snd (put_rev freshl l lbl succs g))

  let insert freshl lbl stmts g =
    let stmt = NodeMap.find lbl g in
    let succs = successors (NodeMap.find lbl g) in
    put_rev freshl (stmt :: List.rev stmts) lbl succs g

  let insert' freshl freshr lbl stmts def g =
    let stmt = NodeMap.find lbl g in
    let succs = successors (NodeMap.find lbl g) in
    let (def, l) = BList.compile freshr def stmts in
    let (lbl, g) = put_rev freshl (stmt :: l) lbl succs g in
    (def, lbl, g)

end

module Trans (Src : GraphType) (Trg : GraphType with type node = Src.node) =
struct
  type node = Src.node

  (* keeping this general branching translation just in case *)
  let translate_general
      (fresh : unit -> Trg.node)
      (f : 'a -> node -> Src.statement ->
       'a * Trg.statement list * Trg.statement list list * node list option)
      (info : 'a)
      (g : Src.t) : 'a * Trg.t =

    let rec put_rev stmts src dests graph =
      match stmts, dests with
        | [], [next] -> Trg.NodeMap.add src (Trg.skip next) graph
        | [last], _ -> Trg.NodeMap.add src (Trg.fill_succs last dests) graph
        | last :: stmts, _ ->
          let new_l = fresh () in
          let graph = Trg.NodeMap.add new_l (Trg.fill_succs last dests) graph in
          put_rev stmts src [new_l] graph
        | _ ->
          invalid_arg "successors of statement and translation do not match" in

    let rec put_block stmts blocks src dests graph =
      match stmts, blocks, dests with
        | [], [b], ds ->
          let rev_b = List.rev b in
          put_rev rev_b src ds graph
        | [], _, _ ->
          invalid_arg "translation: empty preamble with several destinations"
        | _ ->
          let f (lbls, g) blck dst =
            let lbl = fresh () in
            (lbl :: lbls, put_block [] [blck] lbl [dst] g) in
          let (lbls, graph) = List.fold_left2 f ([], graph) blocks dests in
          put_rev (List.rev stmts) src lbls graph in

    let trans lbl stmt (x, graph) =
      let (y, stmts, blocks, redirects) = f x lbl stmt in
      let succs = match redirects with
        | Some x -> x
        | None -> Src.successors stmt in
      (y, put_block stmts blocks lbl succs graph) in

    Src.NodeMap.fold trans g (info, Trg.NodeMap.empty)

  let translate fresh f =
    let f' x lbl stmt =
      let (y, stmts) = f x lbl stmt in
      (y, [], [stmts], None) in
    translate_general fresh f'

  let translate_with_redirects fresh f =
    let f' x lbl stmt =
      let (y, stmts, redirects) = f x lbl stmt in
      (y, [], [stmts], redirects) in
    translate_general fresh f'

  let translate_pure fresh f g =
    let f' () l s = ((), [], [f l s], None) in
    snd (translate_general fresh f' () g)

  let translate_pure_with_redirects fresh f g =
    let f' () l s =
      let (res, redirects) = f l s in
      ((), [], [res], redirects) in
    snd (translate_general fresh f' () g)

  open BList

  let translate_with_redirects' freshl freshr f def g =
    let f' def lbl stmt =
      let (stmts, redirects) = f lbl stmt in
      let (def, stmts) = BList.compile freshr def stmts in
      (def, stmts, redirects) in
    translate_with_redirects freshl f' def g

  let translate' freshl freshr f =
    let f' lbl stmt = (f lbl stmt, None) in
    translate_with_redirects' freshl freshr f'



  (* let translate f = *)
  (*   let f' x _ stmt = *)
  (*     let (y, stmts) = f x stmt in *)
  (*     (y, stmts, [], [Src.successors stmt]) in *)
  (*   translate_general f' *)

  (* let translate_with_redirects f = *)
  (*   let f' x _ stmt = *)
  (*     let (y, stmts, redirects) = f x stmt in *)
  (*     (y, stmts, [], [redirects]) in *)
  (*   translate_general f' *)

  (* let translate_pure f fresh g = *)
  (*   let f' () _ stmt = *)
  (*     ((), f stmt, [], [Src.successors stmt]) in *)
  (*   snd (translate_general f' fresh () g) *)

  (* let translate_pure_with_redirects f fresh g = *)
  (*   let f' () _ stmt = *)
  (*     let (res, redirects) = f stmt in *)
  (*     ((), res, [], [redirects]) in *)
  (*   snd (translate_general f' fresh () g) *)

end