Source file tree.ml

(** A tree widget. Interface is not stable yet. *)

open Tsdl

type 'a tree_spec = {
    is_leaf : 'a -> bool ;
    subs : 'a -> 'a list Lwt.t ;
    can_select : 'a -> bool ;
    can_unselect : 'a -> bool ;
  }

let tree_spec ?(can_select=fun _ -> true)
  ?(can_unselect=fun _ -> true)
    ~is_leaf ~subs () =
    { is_leaf ; subs ;
      can_select ; can_unselect ;
    }

type ('a, 'b) node_label =
  { create : Canvas.group -> 'b ;
    update : selected:bool -> expanded:bool -> 'a -> 'b -> unit ;
    remove : Canvas.group -> 'b -> unit ;
  }

let node_label ~create ~update ~remove = { create ; update ; remove }

let text_collapsed = "▶"
let text_expanded = "▼"
module Id = Misc.Id()
module IMap = Map.Make(Id)
module ISet = Set.Make(Id)

let key_expand = Sdl.K.kp_plus
let key_collapse = Sdl.K.kp_minus
let key_select = Sdl.K.return
let key_next = Sdl.K.down
let key_prev = Sdl.K.up
let key_top = Sdl.K.home
let key_bottom = Sdl.K.kend

type ('a, 'b) node = {
    id: Id.t ;
    parent: ('a, 'b) node option;
    mutable data : 'a ;
    group : Canvas.group ;
    fold_btn : Canvas.label option ;
    label_group : Canvas.group ;
    label : 'b ;
    mutable selected : bool ;
    mutable sub_group : Canvas.group option ;
    mutable expanded: bool ;
    mutable children : ('a, 'b) node list ;
  }

type _ Events.ev +=
| Node_collapsed : (('a * ('a,'b) node) -> bool) Events.ev
| Node_expanded : (('a * ('a,'b) node) -> bool) Events.ev
| Node_clicked : ((Widget.button_ev * 'a * ('a,'b) node) -> bool) Events.ev
| Node_selected : (('a * ('a,'b) node) -> bool) Events.ev
| Node_unselected : (('a * ('a,'b) node) -> bool) Events.ev

class ['a, 'b] tree ?(class_="tree") ?name ?props (spec: 'a tree_spec) (label : ('a,'b) node_label) =
  object(self)
    inherit Canvas.canvas ~class_ ?name ?props () as super
    val mutable text_collapsed_width = 0
    val nodes_mutex = Lwt_mutex.create ()
    val mutable nodes : ('a,'b) node IMap.t = IMap.empty
    val mutable roots = ([] : ('a,'b) node list)
    val mutable selection = ISet.empty

    method selection_mode = self#get_p Props.selection_mode
    method set_selection_mode = self#set_p Props.selection_mode

    method selected_nodes =
      let l = ISet.elements selection in
      List.fold_left
        (fun acc id ->
           match IMap.find_opt id nodes with
           | None -> acc
           | Some n -> n :: acc)
        [] l

    method private connect_node_keys node =
      let _ = node.group#connect Widget.Key_pressed
        (fun ev ->
           Log.debug (fun m -> m "%s key pressed = %s" node.group#me
             (Sdl.(get_key_name ev.Widget.key)));
           match ev.Widget.key with
           | k when k = key_expand -> self#expand node; true
           | k when k = key_collapse -> self#collapse node ; true
           | k when k = key_select -> self#select_or_unselect_node node
           | k when k = key_next -> node.label_group#focus_next
           | k when k = key_prev -> node.label_group#focus_prev
           | k when k = key_top ->
               (match roots with
                | [] -> false
                | n :: _ -> self#node_grab_focus n
               )
           | k when k = key_bottom ->
               (match List.rev roots with
                | [] -> false
                | n :: _ -> self#node_grab_focus n
               )
          | _ -> false
        )
      in
      ()

    method node_grab_focus (node :('a,'b) node) = node.group#grab_focus()

    method private mk_node ?parent t group =
      let group = Canvas.group ~group () in
      let label_group = Canvas.group ~class_:(class_^"_label") ~group () in
      let nlabel = label.create label_group in
      let fold_btn =
        match spec.is_leaf t with
        | true -> None
        | false ->
            let b = Canvas.label ~group text_collapsed in
            Some b
      in
      let node = {
          id = Id.gen() ;
          parent ; data = t ; group ; label_group ;
          sub_group = None ; selected = false ;
          fold_btn ; label = nlabel; expanded = false ; children = [] ;
        }
      in
      self#update_node_display node ;
      let _ = label_group#connect Widget.Clicked (self#on_label_clicked node) in
      self#connect_node_keys node ;
      let () = match fold_btn with
        | None -> ()
        | Some b ->
            let _ = b#connect Widget.Clicked
              (fun ev ->
                 if ev.Widget.button = 1 then
                   self#switch_state node
                 else
                   false)
            in
            ()
      in
      nodes <- IMap.add node.id node nodes;
      node

    method add_root t =
      roots <- roots @ [t] ;
      self#need_resize

    method set_roots l =
      List.iter root#remove_item root#items ;
      nodes <- IMap.empty ;
      let l = List.map (fun r -> self#mk_node r root) l in
      roots <- l;
      self#update_coords

    method private update_coords =
      let rec iter_node y n =
        n.label_group#move ~x:text_collapsed_width ();
        n.group#move ~y ();
        let h =
          match n.expanded with
          | false -> (n.label_group#geometry).G.h
          | true ->
              (match n.sub_group with
               | None -> ()
               | Some g ->
                   g#move ~y: n.label_group#geometry.G.h ()
              );
              let _ = iter 0 n.children in
              (n.group#geometry).G.h
        in
        let h = match n.fold_btn with
        | None -> h
        | Some btn -> max btn#geometry.G.h h
        in
        y + h

      and iter y = function
      | [] -> 0
      | n :: q ->
          let y = iter_node y n in
          iter y q
      in
      ignore(iter 0 roots);
      self#need_resize

    method switch_state node =
      if node.expanded then
        self#collapse node
      else
        self#expand node;
      true

    method collapse node =
      if node.expanded then
        (
         node.expanded <- false ;
         let () =
           match node.fold_btn with
           | None -> ()
           | Some b -> b#set_text text_collapsed
         in
         let () = match node.sub_group with
           | None -> ()
           | Some g ->
               g#set_visible true ;
               node.group#remove_item g#as_full_item
         in
         self#update_coords ;
         self#trigger_event_unit Node_collapsed (node.data, node)
        )
      else
        ()

    method expand (node:('a,'b) node) =
      if not node.expanded then
        (
         match node.fold_btn with
         | None -> ()
         | Some b ->
             node.expanded <- true ;
             self#freeze ;
             b#set_text text_expanded;
             let f () =
               let%lwt () = match node.sub_group with
                 | None ->
                     let%lwt subs = spec.subs node.data in
                     let g = Canvas.group ~x:text_collapsed_width ~group:node.group () in
                     let children = List.map
                       (fun t -> self#mk_node ~parent:node t g)
                         subs
                     in
                     node.children <- children ;
                     node.sub_group <- Some g;
                     Lwt.return_unit
                 | Some g ->
                     Log.info (fun m -> m "node.group%s#geometry: %a" node.group#me G.pp node.group#geometry);
                     node.group#add_item g#as_full_item;
                     Log.info (fun m -> m "subgroup %s added back to %s with geometry %a"
                        g#me node.group#me G.pp g#geometry);
                     Log.info (fun m -> m "node.group%s#geometry: %a" node.group#me G.pp node.group#geometry);
                     g#set_visible true ;
                     Lwt.return_unit
               in
               self#update_coords;
               self#unfreeze ;
               self#trigger_event_unit Node_expanded (node.data, node);
               Lwt.return_unit
             in
             Lwt.async f
        )
      else
        ()

    method unselect_all =
      match self#selected_nodes with
      | [] -> true
      | selected_nodes ->
          match
            List.for_all (fun n -> spec.can_unselect n.data)
              selected_nodes
          with
          | false -> false
          | true ->
              selection <- ISet.empty ;
              List.iter
                (fun n ->
                   n.selected <- false;
                   self#update_node_display n;
                   self#trigger_event_unit Node_unselected (n.data, n))
                  selected_nodes;
              true

    method select_node node =
      if ISet.mem node.id selection then
        true
      else
        match spec.can_select node.data with
        | false -> false
        | true ->
            let selected =
              match self#selection_mode with
              | Props.Sel_none -> false
              | Sel_single
              | Sel_browse ->
                  if self#unselect_all then
                    (selection <- ISet.singleton node.id;
                     true
                    )
                  else
                    false
              | Sel_multiple ->
                  selection <- ISet.add node.id selection ;
                  true
            in
            if selected then
              (
               node.selected <- true;
               self#update_node_display node ;
               self#trigger_event_unit Node_selected (node.data, node) ;
               true
              )
            else
              false

    method unselect_node node =
      match self#selection_mode with
      | Props.Sel_browse -> false
      | _ ->
          if ISet.mem node.id selection then
            match spec.can_unselect node.data with
            | false -> false
            | true ->
                selection <- ISet.remove node.id selection ;
                node.selected <- false;
                self#update_node_display node;
                self#trigger_event_unit Node_unselected (node.data, node);
                true
          else
            true

    method select_or_unselect_node node =
      if ISet.mem node.id selection then
        if Key.is_mod_pressed Tsdl.Sdl.Kmod.ctrl then
          self#unselect_node node
        else
          true
      else
        match spec.can_select node.data with
        | false -> false
        | true ->
            let b =
              if Key.is_mod_pressed Tsdl.Sdl.Kmod.ctrl then
                true
              else
                self#unselect_all
            in
            match b with
            | false -> false
            | true -> self#select_node node

    method on_label_clicked node ev =
      match self#trigger_event Node_clicked (ev, node.data, node) with
      | true -> true
      | false ->
          (* TODO: handle multiselection using SHIFT keys *)
          if ev.Widget.button = 1 then
            (
             self#select_or_unselect_node node ;
             true
            )
          else
            false

    method private rec_remove_node (node: ('a,'b) node) =
      let rec iter n =
        match IMap.find_opt node.id nodes with
        | None -> ()
        | Some node ->
            List.iter iter n.children ;
            let () =
              let rec iter n =
                if ISet.mem n.id selection then
                  (
                   selection <- ISet.remove n.id selection;
                   self#trigger_event_unit Node_unselected (n.data, n)
                  )
                else
                  ();
                List.iter iter n.children
              in
              iter node
            in
            let pred_out n = not (Id.equal n.id node.id) in
            let filter_out = List.filter pred_out in
            let parent_group =
              match node.parent with
              | None -> roots <- filter_out roots; Some root
              | Some p ->
                  p.children <- filter_out p.children;
                  p.sub_group
            in
            Option.iter (fun g -> g#remove_item node.group#as_full_item) parent_group ;
            label.remove node.label_group node.label ;
      in
      iter node

    method remove_node (node: ('a,'b) node) =
      self#rec_remove_node node ;
      self#update_coords

    method remove_nodes nodes =
      List.iter self#rec_remove_node nodes ;
      self#update_coords

    method node_label (node: ('a,'b) node) = node.label

    method private update_node_display node =
      label.update ~selected:node.selected ~expanded:node.expanded
          node.data node.label;
      node.label_group#set_selected node.selected

    method update_node ?(with_subs=false) node data =
      node.data <- data ;
      self#update_node_display node;
      match node.sub_group with
      | Some g when with_subs ->
          self#remove_nodes node.children ;
          if node.expanded then
            (
             node.group#remove_item g#as_full_item;
             node.sub_group <- None ;
             node.expanded <- false ;
             self#expand node
            )
          else
            ()
      | _ -> self#update_coords

    method connect_node_collapsed (f:('a * ('a,'b) node) -> bool) =
      self#connect Node_collapsed f
    method connect_node_expanded (f:('a * ('a,'b) node) -> bool) =
      self#connect Node_expanded f
    method connect_node_clicked (f:(Widget.button_ev * 'a * ('a,'b) node) -> bool) =
      self#connect Node_clicked f
    method connect_node_selected (f:('a * ('a,'b) node) -> bool) =
      self#connect Node_selected f
    method connect_node_unselected (f:('a * ('a,'b) node) -> bool) =
      self#connect Node_unselected f

    initializer
      let indent =
        let fd = Props.(get props font_desc) in
        fd.Font.size
      in
      text_collapsed_width <- indent;

  end

let tree ?class_ ?name ?props ?pack ?selection_mode spec label =
  let w = new tree ?class_ ?name ?props spec label in
  Option.iter w#set_selection_mode selection_mode ;
  Widget.may_pack ?pack w#coerce ;
  w