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

open Format
open Ata

type jump =
  | FIRST_CHILD of StateSet.t
  | NEXT_SIBLING of StateSet.t
  | FIRST_ELEMENT of StateSet.t
  | NEXT_ELEMENT of StateSet.t
  | TAGGED_DESCENDANT of StateSet.t * Tag.t
  | TAGGED_FOLLOWING of StateSet.t * Tag.t
  | SELECT_DESCENDANT of StateSet.t * Ptset.Int.t * Tree.unordered_set
  | SELECT_FOLLOWING of StateSet.t * Ptset.Int.t * Tree.unordered_set
  | TAGGED_CHILD of StateSet.t * Tag.t
  | TAGGED_FOLLOWING_SIBLING of StateSet.t * Tag.t
  | SELECT_CHILD of StateSet.t * Ptset.Int.t * Tree.unordered_set
  | SELECT_FOLLOWING_SIBLING of StateSet.t * Ptset.Int.t * Tree.unordered_set
  | TAGGED_SUBTREE of StateSet.t * Tag.t
  | ELEMENT_SUBTREE of StateSet.t

type dir = DIR_LEFT | DIR_RIGHT
let _nop = None
let _first_child s = Some (FIRST_CHILD s)
let _next_sibling s = Some (NEXT_SIBLING s)
let _first_element s = Some (FIRST_ELEMENT s)
let _next_element s = Some (NEXT_ELEMENT s)
let _tagged_descendant s t = Some (TAGGED_DESCENDANT(s,t))
let _tagged_following s t = Some (TAGGED_FOLLOWING(s,t))
let _select_descendant s t = Some (SELECT_DESCENDANT(s,t, Tree.unordered_set_of_set t))
let _select_following s t = Some (SELECT_FOLLOWING(s,t, Tree.unordered_set_of_set t))
let _tagged_child s t = Some (TAGGED_CHILD(s,t))
let _tagged_following_sibling s t = Some (TAGGED_FOLLOWING_SIBLING(s,t))
let _select_child s t = Some (SELECT_CHILD(s,t, Tree.unordered_set_of_set t))
let _select_following_sibling s t = Some (SELECT_FOLLOWING_SIBLING(s,t, Tree.unordered_set_of_set t))
let _tagged_subtree s t = Some (TAGGED_SUBTREE (s, t))
let _element_subtree s = Some (ELEMENT_SUBTREE s)


let jump_stat_table = Hashtbl.create 17
let jump_stat_init () = Hashtbl.clear jump_stat_table
let jump_stat j =
  let i = try Hashtbl.find jump_stat_table j with Not_found -> 0 in
    Hashtbl.replace jump_stat_table j (i+1)

let print_jump fmt j =
  match j with
    | FIRST_CHILD _ -> fprintf fmt "first_child"
    | NEXT_SIBLING _ -> fprintf fmt "next_sibling"
    | FIRST_ELEMENT _ -> fprintf fmt "first_element"
    | NEXT_ELEMENT _ -> fprintf fmt "next_element"

    | TAGGED_DESCENDANT (_, tag) -> fprintf fmt "tagged_descendant(%s)" (Tag.to_string tag)

    | TAGGED_FOLLOWING (_, tag) -> fprintf fmt "tagged_following(%s)" (Tag.to_string tag)

    | SELECT_DESCENDANT (_, tags, _) -> fprintf fmt "select_descendant(%a)"
	TagSet.print (TagSet.inj_positive tags)

    | SELECT_FOLLOWING (_, tags, _) -> fprintf fmt "select_following(%a)"
	TagSet.print (TagSet.inj_positive tags)

    | TAGGED_CHILD (_, tag) -> fprintf fmt "tagged_child(%s)" (Tag.to_string tag)

    | TAGGED_FOLLOWING_SIBLING (_, tag) ->
	fprintf fmt "tagged_following_sibling(%s)" (Tag.to_string tag)

    | SELECT_CHILD (_, tags, _) -> fprintf fmt "select_child(%a)"
	TagSet.print (TagSet.inj_positive tags)

    | SELECT_FOLLOWING_SIBLING (_, tags, _) -> fprintf fmt "select_following_sibling(%a)"
	TagSet.print (TagSet.inj_positive tags)

    | TAGGED_SUBTREE (_, tag) -> fprintf fmt "tagged_subtree(%s)" (Tag.to_string tag)
    | ELEMENT_SUBTREE (_) -> fprintf fmt "element_subtree"

let jump_stat_summary fmt =
  fprintf fmt "Jump function summary:\n%!";
  Hashtbl.iter (fun k v -> fprintf fmt "%i calls to %a\n" v print_jump k) jump_stat_table;
  fprintf fmt "%!"


type opcode =
  | CACHE
  | RETURN
  | LEFT of Translist.t * jump
  | RIGHT of Translist.t * jump
  | BOTH of Translist.t * jump * jump

type t = opcode Cache.Lvl2.t

let dummy = CACHE
let print_opcode fmt o = match o with
  | CACHE  -> fprintf fmt "CACHE"
  | RETURN  -> fprintf fmt "RETURN"
  | LEFT (tl, j) -> fprintf fmt "LEFT(\n[%a], %a)" Translist.print tl print_jump j
  | RIGHT (tl, j) -> fprintf fmt "RIGHT(\n[%a], %a)" Translist.print tl print_jump j
  | BOTH (tl, j1, j2) -> fprintf fmt "BOTH(\n[%a], %a, %a)" Translist.print tl print_jump j1 print_jump j2

let show_stats a =
  let count = ref 0 in
  Cache.Lvl2.iteri (fun _ _ _ b -> if not b then incr count) a;
  eprintf "%!L2JIT: %i used entries\n%!" !count

let create () =
  let v = Cache.Lvl2.create 4096 dummy in
  if !Options.verbose then
    at_exit (fun () -> show_stats v);
  v

let find t tag set = Cache.Lvl2.find t (Uid.to_int set.StateSet.Node.id) tag

let add t tag set v = Cache.Lvl2.add t (Uid.to_int set.StateSet.Node.id) tag v

let collect_trans tag ((a_t, a_s1, a_s2) as acc) (labels, tr) =
  if TagSet.mem tag labels
  then
    let _, _, _, f = Transition.node tr in
    let  s1,  s2 = Formula.st f in
      (Translist.cons tr a_t,
       StateSet.union s1 a_s1,
       StateSet.union s2 a_s2)
  else acc

let has_text l = Ptset.Int.mem Tag.pcdata l

let rec translate_jump tree tag (jkind:Ata.jump_kind) dir s =
  let child, desc, sib, fol = Tree.tags tree tag in
     match jkind, dir with
      | NIL, _ -> None
      | NODE, DIR_LEFT -> Some (FIRST_CHILD s)
      | STAR, DIR_LEFT -> Some (FIRST_ELEMENT s)
      | NODE, DIR_RIGHT -> Some (NEXT_SIBLING s)
      | STAR, DIR_RIGHT -> Some (NEXT_ELEMENT s)
      | JUMP_ONE t, _ ->
	  let l_one, l_many, tagged_one, select_one, any, any_notext =
	    if dir = DIR_LEFT then
	      child, desc, _tagged_child, _select_child,_first_child, _first_element
	    else
	      sib, fol, _tagged_following_sibling, _select_following_sibling,
	    _next_sibling, _next_element
	  in
	  let labels = Ptset.Int.inter l_one t in
	  let c = Ptset.Int.cardinal labels in
	    if c == 0 then None
	    else if Ptset.Int.for_all (fun lab -> not (Ptset.Int.mem lab l_many)) labels then
	      translate_jump tree tag (JUMP_MANY(labels)) dir s
	    else if c == 1 then tagged_one s (Ptset.Int.choose labels)
	    else if c > 5 then if has_text labels then any s else any_notext s
	    else select_one s labels

      | JUMP_MANY t, _ ->
	  let l_many, tagged_many, select_many, any, any_notext =
	    if dir == DIR_LEFT then
	      desc, _tagged_descendant, _select_descendant,_first_child, _first_element
	    else
	      fol, _tagged_following, _select_following, _next_sibling, _next_element
	  in
	  let labels = Ptset.Int.inter l_many t in
	  let c = Ptset.Int.cardinal labels in
	    if c == 0 then _nop
	    else if c == 1 then tagged_many s (Ptset.Int.choose labels)
	    else if c > 5 then if has_text labels then any s else any_notext s
	    else select_many s labels

      | CAPTURE_MANY (t), DIR_LEFT ->
	  if Ptset.Int.is_singleton t then Some (TAGGED_SUBTREE(s, Ptset.Int.choose t))
	  else if t == Tree.element_tags tree then Some (ELEMENT_SUBTREE s)
	  else assert false
      | _ -> assert false

let compute_jump auto tree tag states dir =
  if !Options.no_jump then
    if dir == DIR_LEFT then Some (FIRST_CHILD states)
    else Some (NEXT_SIBLING states)
  else
    let jkind = Ata.top_down_approx auto states tree in
    translate_jump tree tag jkind dir states

let mk_left tr_list j =
  match j with
    Some x -> LEFT(tr_list, x)
  | _ -> RETURN

let mk_right tr_list j =
  match j with
    Some x -> RIGHT(tr_list, x)
  | _ -> RETURN

let mk_both tr_list j1 j2 =
  match j1, j2 with
  | Some x1, Some x2 -> BOTH(tr_list, x1, x2)
  | None, Some x -> RIGHT(tr_list,x)
  | Some x, None -> LEFT(tr_list, x)
  | None, None -> RETURN

let compile cache2 auto tree tag states =
  let tr_list, states1, states2 =
    StateSet.fold
      (fun q acc ->
	 List.fold_left (collect_trans tag)
	   acc
	   (Hashtbl.find auto.trans q))
      states
      (Translist.nil, StateSet.empty, StateSet.empty)
  in
  let op =
    let empty_s1 = StateSet.is_empty states1 in
    let empty_s2 = StateSet.is_empty states2 in
      if empty_s1 && empty_s2 then RETURN
      else if empty_s1 then
	mk_right tr_list
	  (compute_jump auto tree tag states2 DIR_RIGHT)
      else if empty_s2 then
	mk_left tr_list
	  (compute_jump auto tree tag states1 DIR_LEFT)
      else
	let j1 = compute_jump auto tree tag states1 DIR_LEFT in
	let j2 = compute_jump auto tree tag states2 DIR_RIGHT in
	mk_both tr_list j1 j2
  in
    add cache2 tag states op;
    op

let get_transitions = function
  | CACHE  | RETURN  -> failwith "get_transitions"
  | LEFT (tr, _)
  | RIGHT (tr, _)
  | BOTH (tr, _, _) -> tr