section \<open>Visualization of HOL trees\<close>
theory Tree_Visualization
imports
  Main
keywords
  "visualize_tree" :: diag
begin

subsection \<open>Tree visualization infrastructure\<close>

ML \<open>
signature TREE_VISUALIZATION =
sig

type treevis_mode = string

datatype 'a tree = Node of 'a * 'a tree list
datatype 'a hol_node = Real_Node of 'a | Subtree of term
val map_tree : ('a -> 'b) -> 'a tree -> 'b tree

val tree_to_ascii : string tree -> string

type dot_node = {attribs : (string * string) list, compass : string option}
val tree_to_dot : ('a -> dot_node) -> 'a tree -> string
val tree_to_graph : ('a -> Graph_Display.node) -> 'a tree -> Graph_Display.entry list

val visualize_tree : Proof.context -> string option -> term -> string
val visualize_tree_command : Proof.context -> string option -> string -> unit

val treevis_mode: string Config.T
val register_tree_renderer :
  treevis_mode * binding * (Proof.context -> term -> string) -> Context.generic -> Context.generic
val get_all_tree_renderers :
  Context.generic -> (treevis_mode * (string * (Proof.context -> term -> string)) list) list
val get_tree_renderers :
  Context.generic -> treevis_mode -> (string * (Proof.context -> term -> string)) list

end

structure Tree_Visualization : TREE_VISUALIZATION =
struct

type treevis_mode = string

datatype 'a tree = Node of 'a * 'a tree list

fun map_tree f (Node (x, ts)) = Node (f x, map (map_tree f) ts)

fun intercalate _ [] = ""
  | intercalate _ [x] = x
  | intercalate s (x :: xs) = x ^ s ^ intercalate s xs

val unlines = intercalate "\n"

val tree_to_ascii =
  let
    fun draw (Node (x, ts0)) =
      let
        fun shift first other xs = map op^ ((first :: replicate (length xs - 1) other) ~~ xs)
        fun go [] = []
          | go [t] = "|" :: shift "`-" "  " (draw t)
          | go (t :: ts) = "|" :: shift "+-" "| " (draw t) @ go ts
      in
        x :: go ts0
      end
  in
    draw #> unlines
  end

type dot_node = {attribs : (string * string) list, compass : string option}

fun tree_to_dot mk_node t =
  let
    fun go (Node (x, ts)) (acc, compasses, n) =
      let
        val node = Int.toString n
        val {attribs, compass} = mk_node x
        val attribs = map (fn (a, b) => a ^ "=\"" ^ b ^ "\"") attribs
        val compasses =
          case compass of
            NONE => compasses
          | SOME c => (n, c) :: compasses
        val s = node ^ " [" ^ intercalate ", " attribs ^ "]"
        val acc' = acc @ [s]
        fun go' t (acc, compasses, n') =
          let
            val (acc, compasses, n'') = go t (acc, compasses, n')
            val compass =
              case AList.lookup op= compasses n' of
                NONE => ""
              | SOME c => ":" ^ c
            val acc = acc @ [node ^ " -> " ^ Int.toString n' ^ compass ^ ";"]
          in
            (acc, compasses, n'')
          end
      in
        fold go' ts (acc', compasses, n + 1)
      end
    fun indent s = "  " ^ s
    val xs = go t ([], [], 0) |> #1 |> map indent
    val xs = "digraph tree {" :: indent "splines = \"line\"" :: xs @ ["}"]
  in
    unlines xs
  end

fun tree_to_graph mk_node t =
  let
    fun go parents (Node (x, ts)) (acc, n) =
      let
        val node = Int.toString n
        val acc = acc @ [((node, mk_node x), parents)]
      in
        fold (go [node]) ts (acc, n + 1)
      end
    val xs = go [] t ([], 0) |> #1
  in
    xs
  end
  

datatype 'a hol_node = Real_Node of 'a | Subtree of term

val treevis_mode = Attrib.setup_config_string @{binding "treevis_mode"} (K "ascii");

structure Data = Generic_Data
(
  type T = (treevis_mode * ((Proof.context -> term -> string) Name_Space.table)) list;
  val empty : T = [];
  val extend = I;
  fun merge t12 : T = AList.join op= (K Name_Space.merge_tables) t12;
);

fun register_tree_renderer (mode, name, r) ctxt =
  let
    val xs = Data.get ctxt
    val tbl =
      case AList.lookup op= xs mode of
        NONE => Name_Space.empty_table ("tree renderers for " ^ mode ^ " mode")
      | SOME tbl => tbl
    val tbl' = Name_Space.define ctxt false (name, r) tbl |> snd
  in
    Data.put (AList.update op= (mode, tbl') xs) ctxt
  end
  
fun get_all_tree_renderers ctxt = 
  map (apsnd (fn tbl => Name_Space.fold_table cons tbl [])) (Data.get ctxt)

fun get_tree_renderers ctxt mode =
  case AList.lookup op= (Data.get ctxt) mode of
    NONE => []
  | SOME tbl => Name_Space.fold_table cons tbl []

fun gen_visualize_tree prep_term output ctxt evaluator t =
  let
    val mode = Config.get ctxt treevis_mode
    val renderers = get_tree_renderers (Context.Proof ctxt) mode
    val t = prep_term ctxt t
    val t =
      case evaluator of
        SOME "no_eval" => t
      | NONE => Value_Command.value ctxt t
      | SOME e => Value_Command.value_select e ctxt t
  in
    case get_first (fn (_, f) => try (f ctxt) t) renderers of
      NONE => raise TERM ("visualize_tree_command", [t])
    | SOME s => output s
  end

val visualize_tree = gen_visualize_tree Syntax.check_term I
val visualize_tree_command = gen_visualize_tree Syntax.read_term writeln

val opt_evaluator =
  Scan.option (@{keyword "["} |-- Parse.name --| @{keyword "]"})

val opt_mode = 
  Scan.option (Args.parens Parse.name) >>
    (fn x =>
       case x of
         NONE => I
       | SOME mode => Config.put treevis_mode mode)

val _ =
  Outer_Syntax.command @{command_keyword visualize_tree}
    "Print a visual representation of a HOL tree"
    (opt_evaluator -- opt_mode -- Parse.term >> (fn ((s, f), t) =>
      Toplevel.keep (fn st => visualize_tree_command (f (Toplevel.context_of st)) s t)));

end
\<close>

subsection \<open>Setup for binary trees\<close>

(* Taken from HOL-Library.Tree from Isabelle-2018 *)
datatype 'a tree = Leaf ("\<langle>\<rangle>") | Node "'a tree" 'a "'a tree" ("(1\<langle>_,/ _,/ _\<rangle>)")

setup \<open>
let
  open Tree_Visualization;
  fun show_graph entries = 
    let
      val _ = 
        warning ("The graph display may display children of a binary tree " ^
          "in the wrong order. Use with caution!")
      val _ = Graph_Display.display_graph entries 
    in
      ""
    end
  
  fun term_to_string ctxt t = Syntax.string_of_term ctxt t |> YXML.parse_body |> XML.content_of
  
  datatype hol_bintree_node = HOL_Leaf | HOL_Node of term
  
  fun hol_bintree_to_tree (Const (@{const_name Leaf}, _)) =
        Node (Real_Node HOL_Leaf, [])
    | hol_bintree_to_tree (Const (@{const_name Node}, _) $ l $ x $ r) =
        Node (Real_Node (HOL_Node x), map hol_bintree_to_tree [l, r])
    | hol_bintree_to_tree t =
        case fastype_of t of
          Type (@{type_name "tree"}, _) => Node (Subtree t, [])
        | _ => raise TERM ("hol_bintree_to_tree", [t])
  
  fun mk_bintree_dot_node _ (Real_Node HOL_Leaf) = 
        {attribs = [("shape", "square"), ("label", ""), ("width", "0.15"), ("height", "0.15"),
          ("fillcolor", "black"), ("style", "filled"), ("fixedsize", "true")], compass = NONE}
    | mk_bintree_dot_node ctxt (Real_Node (HOL_Node t)) = 
        {attribs = [("label", term_to_string ctxt t)],
         compass = NONE}
    | mk_bintree_dot_node ctxt (Subtree t)  =
        {attribs = [("label", term_to_string ctxt t), ("shape", "triangle"), ("height", "1.1")],
         compass = SOME "n"}
  
  fun mk_bintree_ascii_node _ (Real_Node HOL_Leaf) = "[]"
    | mk_bintree_ascii_node ctxt (Real_Node (HOL_Node t)) = "[" ^ Syntax.string_of_term ctxt t ^ "]"
    | mk_bintree_ascii_node ctxt (Subtree t) = " " ^ Syntax.string_of_term ctxt t
  
  fun mk_bintree_graph_node _ (Real_Node HOL_Leaf) =
        Graph_Display.content_node "" []
    | mk_bintree_graph_node ctxt (Real_Node (HOL_Node t)) = 
        Graph_Display.content_node (term_to_string ctxt t) [Syntax.pretty_term ctxt t]
    | mk_bintree_graph_node ctxt (Subtree t) =
        Graph_Display.content_node (term_to_string ctxt t) [Syntax.pretty_term ctxt t]

  val bdg = @{binding "tree_renderer"}
  val r1 = ("ascii", bdg, 
    fn ctxt => tree_to_ascii o map_tree (mk_bintree_ascii_node ctxt) o hol_bintree_to_tree)
  val r2 = ("dot", bdg, 
    fn ctxt => tree_to_dot (mk_bintree_dot_node ctxt) o hol_bintree_to_tree)
  val r3 = ("graphview", bdg, 
    fn ctxt => show_graph o tree_to_graph (mk_bintree_graph_node ctxt) o hol_bintree_to_tree)
in
  Context.theory_map (fold register_tree_renderer [r1, r2, r3])
end\<close>


subsection \<open>Examples\<close>

fun ins where
  "ins k Leaf = Node Leaf k Leaf"
| "ins k (Node l x r) = (if k \<le> x then Node (ins k l) x r else Node l x (ins k r))"


text \<open>Inserting into a binary search tree\<close>

visualize_tree "foldr ins [5,3,7,6,4,2,9,1::nat] Leaf"
visualize_tree (dot) "foldr ins [5,3,7,6,4,2,9,1::nat] Leaf"
visualize_tree (graphview) "foldr ins [5,3,7,6,4,2,9,1::nat] Leaf"


text \<open>Partially evaluated insertion into a BST\<close>

visualize_tree [no_eval] 
  "\<langle>\<langle>\<rangle>, 1::nat, \<langle>\<langle>\<langle>\<rangle>, 2, \<langle>\<langle>\<langle>\<rangle>, 3, \<langle>\<rangle>\<rangle>, 4, ins 8 \<langle>\<langle>\<langle>\<rangle>, 5, \<langle>\<rangle>\<rangle>, 6, \<langle>\<langle>\<rangle>, 7, \<langle>\<rangle>\<rangle>\<rangle>\<rangle>\<rangle>, 9, \<langle>\<rangle>\<rangle>\<rangle>"
visualize_tree [no_eval] (dot) 
  "\<langle>\<langle>\<rangle>, 1::nat, \<langle>\<langle>\<langle>\<rangle>, 2, \<langle>\<langle>\<langle>\<rangle>, 3, \<langle>\<rangle>\<rangle>, 4, ins 8 \<langle>\<langle>\<langle>\<rangle>, 5, \<langle>\<rangle>\<rangle>, 6, \<langle>\<langle>\<rangle>, 7, \<langle>\<rangle>\<rangle>\<rangle>\<rangle>\<rangle>, 9, \<langle>\<rangle>\<rangle>\<rangle>"
visualize_tree [no_eval] (graphview) 
  "\<langle>\<langle>\<rangle>, 1::nat, \<langle>\<langle>\<langle>\<rangle>, 2, \<langle>\<langle>\<langle>\<rangle>, 3, \<langle>\<rangle>\<rangle>, 4, ins 8 \<langle>\<langle>\<langle>\<rangle>, 5, \<langle>\<rangle>\<rangle>, 6, \<langle>\<langle>\<rangle>, 7, \<langle>\<rangle>\<rangle>\<rangle>\<rangle>\<rangle>, 9, \<langle>\<rangle>\<rangle>\<rangle>"


text \<open>Tree with subtrees\<close>

visualize_tree "\<langle>t1, a, \<langle>\<langle>t2, c, t3\<rangle>, b, t4\<rangle>\<rangle>"
visualize_tree (dot) "\<langle>t1, a, \<langle>\<langle>t2, c, t3\<rangle>, b, t4\<rangle>\<rangle>"
visualize_tree (graphview) "\<langle>t1, a, \<langle>\<langle>t2, c, t3\<rangle>, b, t4\<rangle>\<rangle>"

end