Source file ngraph.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2024                                                    *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  You can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It is distributed in the hope that it will be useful,                 *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)
open Base

type t = {
  output : Node.t;
  succs : (int, Node.t list) Base.Hashtbl.t;
}

let output t = t.output

(** TODO: some other invariants must be checked e.g only one input *)
let create output =
  let succs = Base.Hashtbl.create (module Base.Int) in
  let check_node node =
    List.iter
      ~f:(fun p -> Base.Hashtbl.add_multi succs ~key:p.id ~data:node)
      (Node.preds node)
  in
  (* Add the key of the output nodes so that all the nodes are in succs *)
  Base.Hashtbl.add_exn succs ~key:output.Node.id ~data:[];
  Node.iter_rec check_node output;
  { output; succs }

let input_shape g =
  let r = ref None in
  let check_node n =
    match n.Node.descr with Input { shape } -> r := Some shape | _ -> ()
  in
  Node.iter_rec check_node g.output;
  Option.value_exn !r

let succs t node = Base.Hashtbl.find_exn t.succs node.Node.id
let iter_vertex f t = Node.iter_rec f t.output

let iter_succ f t node =
  List.iter ~f (Base.Hashtbl.find_multi t.succs node.Node.id)

let pp fmt t =
  iter_vertex (fun v -> Fmt.pf fmt "@[%a@]@ " Node.pp_descr v.descr) t

let pp_debug fmt t =
  iter_vertex
    (fun v ->
      Fmt.pf fmt "@[%i: %a(%a) : %a@]@ " v.id Node.pp_descr v.descr
        Fmt.(list ~sep:comma (using (fun x -> x.Node.id) int))
        (Node.preds v) Shape.pp v.shape)
    t

let nodes t =
  let l = ref [] in
  iter_vertex (fun v -> l := v :: !l) t;
  !l

module M = Graph.Topological.Make

module GFloat = struct
  type nonrec t = t

  let iter_edges_e f t =
    iter_vertex (fun n -> List.iter ~f:(fun n' -> f (n', n)) (Node.preds n)) t

  module V = Node

  module E = struct
    type t = V.t * V.t

    let src = fst
    let dst = snd
  end

  let iter_vertex = iter_vertex
  let iter_succ = iter_succ
end

module Dot = Graph.Graphviz.Dot (struct
  include GFloat (* use the graph module from above *)

  let node_label (v : Node.t) = Node.show v
  let edge_attributes (_, _) = []
  let default_edge_attributes _ = []
  let get_subgraph _ = None
  let vertex_attributes v = [ `Shape `Box; `Label (node_label v) ]
  let vertex_name (v : Node.t) = Int.to_string v.id
  let default_vertex_attributes _ = []
  let graph_attributes _ = []
end)

let grapheasy g =
  try
    let cin, cout =
      Unix.open_process_args "graph-easy"
        [| "graph-easy"; "--from=graphviz"; "--as=boxart" |]
    in
    Dot.output_graph cout g;
    Stdlib.close_out cout;
    let ascii = Stdio.In_channel.input_all cin in
    ignore (Unix.close_process (cin, cout));
    ascii
  with exn ->
    Fmt.str "Error graph-easy call: %s" (Stdlib.Printexc.to_string exn)