(***********************************************************************)
(*                                                                     *)
(*                               TAToo                                 *)
(*                                                                     *)
(*                     Kim Nguyen, LRI UMR8623                         *)
(*                   Université Paris-Sud & CNRS                       *)
(*                                                                     *)
(*  Copyright 2010-2013 Université Paris-Sud and Centre National de la *)
(*  Recherche Scientifique. All rights reserved.  This file is         *)
(*  distributed under the terms of the GNU Lesser General Public       *)
(*  License, with the special exception on linking described in file   *)
(*  ../LICENSE.                                                        *)
(*                                                                     *)
(***********************************************************************)

INCLUDE "utils.ml"
INCLUDE "debug.ml"

open Format
open Misc

type stats = { run : int;
               tree_size : int;
               fetch_trans_cache_access : int;
               fetch_trans_cache_hit : int;
               eval_trans_cache_access : int;
               eval_trans_cache_hit : int;
             }

let fetch_trans_cache_hit = ref 0
let fetch_trans_cache_access = ref 0
let eval_trans_cache_hit = ref 0
let eval_trans_cache_access = ref 0
let reset_stat_counters () =
  fetch_trans_cache_hit := 0;
  fetch_trans_cache_access := 0;
  eval_trans_cache_hit := 0;
  eval_trans_cache_access := 0


module Make (T : Tree.S) =
 struct

   module NodeSummary =
   struct
     (* Pack into an integer the result of the is_* and has_ predicates
        for a given node *)
     type t = int
     let dummy = -1
    (*
      ...44443210
      ...4444 -> kind
      3 -> has_right
      2 -> has_left
      1 -> is_right
      0 -> is_left
    *)
     let is_left (s : t) : bool =
       s land 1 != 0

     let is_right (s : t) : bool =
       s land 0b10 != 0

     let has_left (s : t) : bool =
       s land 0b100 != 0

     let has_right (s : t) : bool =
       s land 0b1000 != 0

     let kind (s : t) : Tree.NodeKind.t =
       Obj.magic (s lsr 4)

     let make is_left is_right has_left has_right kind =
       (int_of_bool is_left) lor
         ((int_of_bool is_right) lsl 1) lor
         ((int_of_bool has_left) lsl 2) lor
         ((int_of_bool has_right) lsl 3) lor
         ((Obj.magic kind) lsl 4)
   end

   let dummy_set = StateSet.singleton State.dummy
   open Bigarray
   type run = {
     tree : T.t ;
     (* The argument of the run *)
     auto : Ata.t;
     (* The automaton to be run *)
     sat: StateSet.t array;
     (* A mapping from node preorders to states satisfied at that node *)
     mutable pass : int;
     (* Number of run we have performed *)
     mutable fetch_trans_cache : Ata.Formula.t Cache.N2.t;
     (* A cache from states * label to list of transitions *)
     mutable td_cache : StateSet.t Cache.N6.t;
     mutable bu_cache : StateSet.t Cache.N6.t;
     (* Two 6-way caches used during the top-down and bottom-up phase
        label * self-set * fc-set * ns-set * parent-set * node-shape -> self-set
     *)
     node_summaries: (int, int16_unsigned_elt, c_layout) Array1.t;
   }


   let dummy_form = Ata.Formula.stay State.dummy

   let make auto tree =
     let len = T.size tree in
     {
       tree = tree;
       auto = auto;
       sat = Array.create len StateSet.empty;
       pass = 0;
       fetch_trans_cache = Cache.N2.create dummy_form;
       td_cache = Cache.N6.create dummy_set;
       bu_cache = Cache.N6.create dummy_set;
       node_summaries = let ba = Array1.create int16_unsigned c_layout len in
                        Array1.fill ba 0; ba
     }

   let get_form fetch_trans_cache auto tag q =
     let phi =
       incr fetch_trans_cache_access;
       Cache.N2.find fetch_trans_cache (tag.QName.id :> int) (q :> int)
     in
     if phi == dummy_form then
       let phi = Ata.get_form auto tag q in
       let () =
         Cache.N2.add
           fetch_trans_cache
           (tag.QName.id :> int)
           (q :> int) phi
       in phi
     else begin
       incr fetch_trans_cache_hit;
       phi
     end


   let eval_form phi fcs nss ps ss summary =
     let open Ata in
         let rec loop phi =
           begin match Formula.expr phi with
           | Boolean.False -> false
           | Boolean.True -> true
           | Boolean.Atom (a, b) ->
               begin
                 let open NodeSummary in
                     match a.Atom.node with
                     | Move (m, q) ->
                       b && StateSet.mem q (
                           match m with
                             `First_child -> fcs
                           | `Next_sibling -> nss
                           | `Parent | `Previous_sibling ->  ps
                           | `Stay ->  ss
                       )
                     | Is_first_child -> b == is_left summary
                     | Is_next_sibling -> b == is_right summary
                     | Is k -> b == (k == kind summary)
                     | Has_first_child -> b == has_left summary
                     | Has_next_sibling -> b == has_right summary
               end
           | Boolean.And(phi1, phi2) -> loop phi1 && loop phi2
           | Boolean.Or (phi1, phi2) -> loop phi1 || loop phi2
           end
         in
         loop phi


   let eval_trans_aux auto fetch_trans_cache tag fcs nss ps sat todo summary =
     StateSet.fold (fun q (a_sat) ->
       let phi =
         get_form fetch_trans_cache auto tag q
       in
       if eval_form phi fcs nss ps a_sat summary then
         StateSet.add q a_sat
       else a_sat
     ) todo sat


   let rec eval_trans_fix auto fetch_trans_cache tag fcs nss ps sat todo summary =
     let new_sat =
       eval_trans_aux auto fetch_trans_cache tag fcs nss ps sat todo summary
     in
     if new_sat == sat then sat else
       eval_trans_fix auto fetch_trans_cache tag fcs nss ps new_sat todo summary


   let eval_trans auto fetch_trans_cache eval_cache tag fcs nss ps ss todo summary =
     let fcsid = (fcs.StateSet.id :> int) in
     let nssid = (nss.StateSet.id :> int) in
     let psid = (ps.StateSet.id :> int) in
     let ssid = (ss.StateSet.id :> int) in
     let tagid = (tag.QName.id :> int) in
     let res = Cache.N6.find eval_cache tagid summary ssid fcsid nssid psid in
     incr eval_trans_cache_access;
     if res != dummy_set then begin incr eval_trans_cache_hit; res end
     else let new_sat =
            eval_trans_fix auto fetch_trans_cache tag fcs  nss ps ss todo summary
          in
          Cache.N6.add eval_cache tagid summary ssid fcsid nssid psid new_sat;
          new_sat


   let unsafe_get a i = if i < 0 then StateSet.empty else Array.unsafe_get a i

   let top_down run =
    let i = run.pass in
    let tree = run.tree in
    let auto = run.auto in
    let states_by_rank = Ata.get_states_by_rank auto in
    let td_todo = states_by_rank.(i) in
    let bu_todo = if i + 1 = Array.length states_by_rank then StateSet.empty
      else
        states_by_rank.(i+1)
    in
    let rec loop_td_and_bu node parent parent_sat =
      if node == T.nil then StateSet.empty else begin
        let node_id = T.preorder tree node in
        let fc = T.first_child tree node in
        let ns = T.next_sibling tree node in
        let tag = T.tag tree node in
        (* We enter the node from its parent *)
        let summary =
          let s = Array1.unsafe_get run.node_summaries node_id in
          if  s != 0 then s else
            let s =
              NodeSummary.make
                (node == T.first_child tree parent) (*is_left *)
                (node == T.next_sibling tree parent)(*is_right *)
                (fc != T.nil) (* has_left *)
                (ns != T.nil) (* has_right *)
                (T.kind tree node) (* kind *)
            in
            run.node_summaries.{node_id} <- s; s
        in
        let status0 = unsafe_get run.sat node_id in
        (* get the node_statuses for the first child, next sibling and parent *)
        let fcs = unsafe_get run.sat (T.preorder tree fc) in
        let nss = unsafe_get run.sat (T.preorder tree ns) in
        (* evaluate the transitions with all this statuses *)
        let status1 =
          eval_trans auto run.fetch_trans_cache run.td_cache tag fcs nss
            parent_sat
            status0 td_todo summary
        in
        (* update the cache if the status of the node changed *)
        if status1 != status0 then run.sat.(node_id) <- status1;
        let fcs1 = loop_td_and_bu fc node status1 in
        if bu_todo == StateSet.empty then
          loop_td_and_bu ns node status1 (* tail call *)
        else
          let nss1 = loop_td_and_bu ns node status1 in
          let status2 =
            eval_trans auto run.fetch_trans_cache run.bu_cache tag fcs1 nss1
              parent_sat
              status1 bu_todo summary
          in
          if status2 != status1 then run.sat.(node_id) <- status2;
          status2
      end
    in
    let _ = loop_td_and_bu (T.root tree) T.nil StateSet.empty in
    run.pass <- run.pass + 2


  let get_results run =
    let cache = run.sat in
    let auto = run.auto in
    let tree = run.tree in
    let sel_states = Ata.get_selecting_states auto in
    let rec loop node acc =
      if node == T.nil then acc
      else
        let acc0 = loop (T.next_sibling tree node) acc in
        let acc1 = loop (T.first_child tree node) acc0 in

        if StateSet.intersect
          cache.(T.preorder tree node)(* NodeStatus.node.sat *)
          sel_states then node::acc1
        else acc1
    in
    loop (T.root tree) []


  let get_full_results run =
    let cache = run.sat(*tatus*) in
    let auto = run.auto in
    let tree = run.tree in
    let res_mapper = Hashtbl.create MED_H_SIZE in
    let () =
      StateSet.iter
        (fun q -> Hashtbl.add res_mapper q [])
        (Ata.get_selecting_states auto)
    in
    let dummy = [ T.nil ] in
    let res_mapper = Cache.N1.create dummy in
    let () =
      StateSet.iter
        (fun q -> Cache.N1.add res_mapper (q :> int) [])
        (Ata.get_selecting_states auto)
    in
    let rec loop node =
      if node != T.nil then
        let () = loop (T.next_sibling tree node) in
        let () = loop (T.first_child tree node) in
        StateSet.iter
          (fun q ->
            let res = Cache.N1.find res_mapper (q :> int) in
            if res != dummy then
              Cache.N1.add res_mapper (q :> int) (node::res)
          )
          cache.(T.preorder tree node)(* NodeStatus.node.sat *)
    in
    loop (T.root tree);
    (StateSet.fold_right
       (fun q acc -> (q, Cache.N1.find res_mapper (q :> int))::acc)
       (Ata.get_selecting_states auto) [])


  let prepare_run run list =
    let tree = run.tree in
    let auto = run.auto in
    let sat0 = Ata.get_starting_states auto in
    List.iter (fun node ->
      let node_id = T.preorder tree node in
      run.sat.(node_id) <- sat0) list

  let tree_size = ref 0
  let pass = ref 0
  let compute_run auto tree nodes =
    pass := 0;
    tree_size := T.size tree;
    let run = make auto tree in
    prepare_run run nodes;
    let rank = Ata.get_max_rank auto in
    while run.pass <= rank do
      top_down run;
      run.td_cache <- Cache.N6.create dummy_set;
      run.bu_cache <- Cache.N6.create dummy_set;
    done;
    pass := Ata.get_max_rank auto + 1;

    run

  let full_eval auto tree nodes =
    let r = compute_run auto tree nodes in
    get_full_results r

  let eval auto tree nodes =
    let r = compute_run auto tree nodes in
    get_results r

  let stats () = {
    tree_size = !tree_size;
    run = !pass;
    fetch_trans_cache_access = !fetch_trans_cache_access;
    fetch_trans_cache_hit = !fetch_trans_cache_hit;
    eval_trans_cache_access = !eval_trans_cache_access;
    eval_trans_cache_hit = !eval_trans_cache_hit;
  }

end