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

open Format
open Ata

type jump =
  | NOP of unit
  | 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.tag_list
  | SELECT_FOLLOWING of StateSet.t * Ptset.Int.t * Tree.tag_list
  | TAGGED_CHILD of StateSet.t * Tag.t
  | TAGGED_SIBLING of StateSet.t * Tag.t
  | SELECT_CHILD of StateSet.t * Ptset.Int.t * Tree.tag_list
  | SELECT_SIBLING of StateSet.t * Ptset.Int.t * Tree.tag_list
  | TAGGED_SUBTREE of StateSet.t * Tag.t
  | ELEMENT_SUBTREE of StateSet.t

type dir = DIR_LEFT | DIR_RIGHT

let _nop = NOP ()
let _first_child s = FIRST_CHILD s
let _next_sibling s = NEXT_SIBLING s
let _first_element s = FIRST_ELEMENT s
let _next_element s = NEXT_ELEMENT s
let _tagged_descendant s t = TAGGED_DESCENDANT(s,t)
let _tagged_following s t = TAGGED_FOLLOWING(s,t)
let _select_descendant s t = SELECT_DESCENDANT(s,t, Tree.tag_list_of_set t)
let _select_following s t = SELECT_FOLLOWING(s,t, Tree.tag_list_of_set t)
let _tagged_child s t = TAGGED_CHILD(s,t)
let _tagged_following_sibling s t = TAGGED_SIBLING(s,t)
let _select_child s t = SELECT_CHILD(s,t, Tree.tag_list_of_set t)
let _select_following_sibling s t = SELECT_SIBLING(s,t, Tree.tag_list_of_set t)
let _tagged_subtree s t = TAGGED_SUBTREE (s, t)
let _element_subtree s = 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
  | NOP _ -> fprintf fmt "nop"
  | 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_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_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 =
  | RETURN of unit
  | LEFT of Translist.t * jump
  | RIGHT of Translist.t * jump
  | BOTH of Translist.t * jump * jump
  | CACHE of unit

type t = opcode Cache.Lvl2.t

let dummy = CACHE ()
let return = RETURN ()
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 print_cache fmt d =
  let c = Cache.Lvl2.to_array d in
  Array.iteri begin fun tag a ->
  let tagstr = Tag.to_string tag in
  if a != Cache.Lvl2.dummy_line d && tagstr <> "<INVALID TAG>"
  then begin
  fprintf fmt "Entry %s: \n" tagstr;
  Array.iter (fun o -> if o != dummy then begin
  print_opcode fmt o;
  fprintf fmt "\n%!" end) a;
  fprintf fmt "---------------------------\n%!"
  end
  end c
*)
let create () = Cache.Lvl2.create 512 dummy
(*
  let stats fmt c =
  let d = Cache.Lvl2.to_array c in
  let len = Array.fold_left (fun acc a -> Array.length a + acc) 0 d in
  let lvl1 = Array.fold_left (fun acc a -> if Array.length a == 0 then acc else acc+1) 0 d in
  let lvl2 = Array.fold_left (fun acc a ->
  Array.fold_left (fun acc2 a2 -> if a2 == dummy then acc2 else acc2+1)
  acc a) 0 d
  in
  fprintf fmt "L2JIT Statistics:
  \t%i entries
  \t%i used L1 lines
  \t%i used L2 lines
  \ttable size: %ikb\n"
  len lvl1 lvl2 (Ocaml.size_kb d);
  fprintf fmt "%s" "L2JIT Content:\n";
  print_cache fmt c
*)

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 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
  let not_elements =
    Ptset.Int.add Tag.pcdata
      (Ptset.Int.add Tag.attribute
         (Ptset.Int.add Tag.attribute_data
            (Tree.attribute_tags tree)))
  in
  match jkind, dir with
  | NIL, _ -> _nop
  | NODE, DIR_LEFT -> FIRST_CHILD s
  | STAR, DIR_LEFT -> FIRST_ELEMENT s
  | NODE, DIR_RIGHT -> NEXT_SIBLING s
  | STAR, DIR_RIGHT -> 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 _nop
    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
      let not_t = Ptset.Int.diff l_many labels in
      let () =
        LOG(__ "level2-jit" 3 "Would jump for tag %s to labels %a, not relevant tags: %a"
              (Tag.to_string tag)
              TagSet.print (TagSet.inj_positive labels)
              TagSet.print (TagSet.inj_positive not_t))
      in
      if Ptset.Int.subset not_t not_elements then
        if has_text labels then any s else any_notext s
      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 TAGGED_SUBTREE(s, Ptset.Int.choose t)
    else if t == Tree.element_tags tree then ELEMENT_SUBTREE s
    else assert false
  | _ -> assert false

let count = ref 0
let () = at_exit (fun () -> Logger.verbose Format.err_formatter "Compute jump called %i times\n" !count)
module Memo = Hashtbl.Make(struct
  type t = Tag.t * StateSet.t * dir
  let equal (a,b,c) (d,e,f) = a == d && b == e && c == f
  let hash (a, b, c) = HASHINT3(a, Uid.to_int b.StateSet.Node.id, (Obj.magic c))
end)

let memo = Memo.create 1024
let init () = Memo.clear memo

let compute_jump auto tree tag states dir =
  if !Config.no_jump then
    if dir == DIR_LEFT then FIRST_CHILD states
    else NEXT_SIBLING states
  else
    try
      Memo.find memo (tag, states, dir)
    with
      Not_found -> begin
	incr count;
	let jkind = Ata.top_down_approx auto states tree in
	let jump = translate_jump tree tag jkind dir states in
	LOG(__ "level2-jit" 2
              "Computed jumps for %s %a %s, from %a : %a%!"
              (Tag.to_string tag)
              StateSet.print states
              (if dir == DIR_LEFT then "left" else "right")
              Ata.print_kind jkind
              print_jump jump
	);
	Memo.add memo (tag, states, dir) jump; jump
      end

let compile cache2 auto tree tag states =
  let tr_list, states1, states2 =
    Ata.get_trans (*~attributes:(TagSet.inj_positive (Tree.attribute_tags tree))*) auto states tag
  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
      RIGHT (tr_list,
	     compute_jump auto tree tag states2 DIR_RIGHT)
    else if empty_s2 then
      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
      BOTH (tr_list, j1, j2);
  in
  let op = match op with
    (*BOTH(_, NOP _, NOP _) |  LEFT(_, NOP _) | RIGHT(_, NOP _) -> RETURN() *)
    | BOTH(tr, ((NOP _) as l) , NOP _) -> LEFT (tr, l)
    | BOTH(tr, l, NOP _) -> LEFT (tr, l)
    | BOTH(tr, NOP _, r) -> RIGHT (tr, r)
    | _ -> op
  in
  if not !Config.no_cache then add cache2 tag states op;
  op

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