(******************************************************************************)
(*  SXSI : XPath evaluator                                                    *)
(*  Kim Nguyen (Kim.Nguyen@nicta.com.au)                                      *)
(*  Copyright NICTA 2008                                                      *)
(*  Distributed under the terms of the LGPL (see LICENCE)                     *)
(******************************************************************************)
INCLUDE "utils.ml"


external init_lib : unit -> unit = "caml_init_lib"

exception CPlusPlusError of string

let () = Callback.register_exception "CPlusPlusError" (CPlusPlusError "")

let () =  init_lib ()


type tree
type 'a node = private int
type node_kind = [`Text | `Tree ]

external inode : 'a node -> int = "%identity"  
external nodei : int -> 'a node = "%identity"  
let compare_node x y = (inode x) - (inode y)
let equal_node : 'a node -> 'a node -> bool = (==)

  
external parse_xml_uri : string -> int -> bool -> bool -> tree = "caml_call_shredder_uri"         
external parse_xml_string :  string -> int -> bool -> bool -> tree = "caml_call_shredder_string"
  
external tree_save : tree -> Unix.file_descr -> unit = "caml_xml_tree_save"
external tree_load : Unix.file_descr -> tree = "caml_xml_tree_load"
  
external nullt : unit -> 'a node = "caml_xml_tree_nullt"

let nil : [`Tree ] node = nodei ~-1
let nulldoc : [`Text ] node = nodei ~-1
let root : [`Tree ] node = nodei 0

external text_get_tc_text : tree -> [`Text] node -> string = "caml_text_collection_get_text" 
		
external text_is_empty : tree -> [`Text ] node -> bool = "caml_text_collection_empty_text" 

let text_is_empty t n = (equal_node nulldoc n) || text_is_empty t n

external text_is_contains : tree -> string -> bool = "caml_text_collection_is_contains" 
external text_count_contains : tree -> string -> int = "caml_text_collection_count_contains" 
external text_count : tree -> string -> int = "caml_text_collection_count" 
external text_contains : tree -> string -> [`Text ] node array = "caml_text_collection_contains" 
external text_unsorted_contains : tree -> string -> unit = "caml_text_collection_unsorted_contains"
external text_get_cached_text : tree -> [`Text] node -> string = "caml_text_collection_get_cached_text"
    
external tree_root : tree -> [`Tree] node = "caml_xml_tree_root" 
external tree_subtree_size : tree -> [`Tree] node -> int = "caml_xml_tree_subtree_size"
 
let tree_is_nil x = equal_node x nil

external tree_parent : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_parent" "noalloc"
external tree_parent_doc : tree -> [`Text ] node -> [`Tree ] node = "caml_xml_tree_parent_doc" "noalloc"
(*external tree_prev_doc : tree -> [`Text ] node -> [`Tree ] node = "caml_xml_tree_prev_doc" "noalloc" *)
external tree_first_child : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_first_child" "noalloc"
external tree_first_element : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_first_element" "noalloc"
external tree_tagged_child : tree -> [`Tree] node -> Tag.t -> [`Tree] node = "caml_xml_tree_tagged_child" "noalloc" 
external tree_next_sibling : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_next_sibling"  "noalloc"
external tree_next_element : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_next_element"  "noalloc"
external tree_tagged_sibling : tree -> [`Tree] node -> Tag.t -> [`Tree] node = "caml_xml_tree_tagged_sibling" "noalloc"

external tree_prev_sibling : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_prev_sibling" "noalloc"
external tree_is_leaf : tree -> [`Tree] node -> bool = "caml_xml_tree_is_leaf" "noalloc"
external tree_last_child : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_last_child" "noalloc"
external tree_is_first_child : tree -> [`Tree] node -> bool = "caml_xml_tree_is_first_child" "noalloc"


external tree_tag_id : tree -> [`Tree ] node -> Tag.t = "caml_xml_tree_tag_id"  "noalloc"
    

let tree_is_last t n = equal_node nil (tree_next_sibling t n)
    
(*external tree_prev_text : tree -> [`Tree] node -> [`Text ] node = "caml_xml_tree_prev_text" "noalloc" *)

external tree_my_text : tree -> [`Tree] node -> [`Text ] node = "caml_xml_tree_my_text" "noalloc"
(*external tree_next_text : tree -> [`Tree] node -> [`Text ] node = "caml_xml_tree_next_text" "noalloc" *)
external tree_doc_ids : tree -> [`Tree ] node -> [`Text ] node * [`Text ] node = "caml_xml_tree_doc_ids" 

let text_size tree = inode (snd ( tree_doc_ids tree root ))

let text_get_cached_text t (x:[`Text] node) =
  if x == nulldoc then ""
  else 
     text_get_cached_text t x


external tree_text_xml_id : tree -> [`Text ] node -> int = "caml_xml_tree_text_xml_id" "noalloc"
external tree_node_xml_id : tree -> [`Tree ] node -> int = "caml_xml_tree_node_xml_id" "noalloc"
external tree_is_ancestor : tree -> [`Tree ] node -> [`Tree ] node -> bool = "caml_xml_tree_is_ancestor" "noalloc" 
external tree_tagged_desc : tree -> [`Tree ] node -> Tag.t -> [`Tree ] node = "caml_xml_tree_tagged_desc" "noalloc"
external tree_tagged_foll_below : tree -> [`Tree ] node -> Tag.t -> [`Tree ] node -> [`Tree ] node = "caml_xml_tree_tagged_foll_below" "noalloc"
external tree_subtree_tags : tree -> [`Tree ] node -> Tag.t -> int = "caml_xml_tree_subtree_tags" "noalloc"


type unordered_set
external unordered_set_alloc : int -> unordered_set = "caml_unordered_set_alloc"
external unordered_set_length : unordered_set -> int = "caml_unordered_set_length"
external unordered_set_insert : unordered_set -> int -> unit = "caml_unordered_set_set" "noalloc"

external tree_select_child : tree -> [`Tree ] node -> unordered_set -> [`Tree] node = "caml_xml_tree_select_child" "noalloc"
external tree_select_foll_sibling : tree -> [`Tree ] node -> unordered_set -> [`Tree] node = "caml_xml_tree_select_foll_sibling" "noalloc"
external tree_select_desc : tree -> [`Tree ] node -> unordered_set -> [`Tree] node = "caml_xml_tree_select_desc" "noalloc"
external tree_select_foll_below : tree -> [`Tree ] node -> unordered_set -> [`Tree] node -> [`Tree] node = "caml_xml_tree_select_foll_below" "noalloc"


module HPtset = Hashtbl.Make(Ptset.Int)

let vector_htbl = HPtset.create MED_H_SIZE

let ptset_to_vector s =
  try 
    HPtset.find vector_htbl s
  with
      Not_found ->
	let v = unordered_set_alloc (Ptset.Int.cardinal s) in
	let _ = Ptset.Int.iter (fun e -> unordered_set_insert v e) s in
	  HPtset.add vector_htbl s v; v

      
type t = { 
  doc : tree;	          
  ttable : (Tag.t,(Ptset.Int.t*Ptset.Int.t*Ptset.Int.t*Ptset.Int.t)) Hashtbl.t;
}
let subtree_size t i = tree_subtree_size t.doc i
let text_size t = text_size t.doc

module MemUnion = Hashtbl.Make (struct 
      type t = Ptset.Int.t*Ptset.Int.t
      let equal (x,y) (z,t) = (Ptset.Int.equal x z)&&(Ptset.Int.equal y t)
      let equal a b = equal a b || equal b a
      let hash (x,y) =   (* commutative hash *)
	let x = Ptset.Int.hash x 
	and y = Ptset.Int.hash y 
	in
	  if x < y then HASHINT2(x,y) else HASHINT2(y,x)
    end)

module MemAdd = Hashtbl.Make (
  struct 
    type t = Tag.t*Ptset.Int.t
    let equal (x,y) (z,t) = (x == z)&&(Ptset.Int.equal y t)
    let hash (x,y) =  HASHINT2(x,Ptset.Int.hash y)
  end)

let collect_tags tree =
  let h_union = MemUnion.create BIG_H_SIZE in
  let pt_cup s1 s2 =
      try
	MemUnion.find h_union (s1,s2)
      with
	| Not_found -> let s = Ptset.Int.union s1 s2
	  in
	    MemUnion.add h_union (s1,s2) s;s
  in    
  let h_add = MemAdd.create BIG_H_SIZE in
  let pt_add t s =  
    try MemAdd.find h_add (t,s)
    with
      | Not_found -> let r = Ptset.Int.add t s in
	  MemAdd.add h_add (t,s) r;r
  in
  let h = Hashtbl.create BIG_H_SIZE in
  let update t sc sb ss sa =
    let schild,sbelow,ssibling,safter =  
      try
	Hashtbl.find h t 
      with
	| Not_found -> 
	    (Ptset.Int.empty,Ptset.Int.empty,Ptset.Int.empty,Ptset.Int.empty)
    in
      Hashtbl.replace h t 
	(pt_cup sc schild,pt_cup sbelow sb, pt_cup ssibling ss, pt_cup safter sa)
  in
  let rec loop_right id acc_after = 
    if  id == nil
    then Ptset.Int.empty,Ptset.Int.empty,acc_after
    else
    let sibling2,desc2,after2 = loop_right (tree_next_sibling tree id) acc_after in
    let child1,desc1,after1   = loop_left (tree_first_child tree id) after2  in
    let tag = tree_tag_id tree id in
    update tag child1 desc1 sibling2 after2;
    ( pt_add tag sibling2, 
      pt_add tag (pt_cup desc1 desc2),
      pt_cup after1 (pt_cup desc1 desc2) )
  and loop_left id acc_after = 
    if  id == nil
    then Ptset.Int.empty,Ptset.Int.empty,acc_after
    else
    let sibling2,desc2,after2 = loop_right (tree_next_sibling tree id) acc_after in
    let child1,desc1,after1   = loop_left (tree_first_child tree id) after2  in
    let tag = tree_tag_id tree id in
    update tag child1 desc1 sibling2 after2;
    (pt_add tag sibling2, 
     pt_add tag (pt_cup desc1 desc2),
     acc_after )
  in
  let _ = loop_left (tree_root tree) Ptset.Int.empty in h
			  
			  
    

let contains_array = ref [| |]
let contains_index = Hashtbl.create 4096 
let in_array _ i =
  try
    Hashtbl.find contains_index i
  with
      Not_found -> false

let init_contains t s = 
  let a = text_contains t.doc s 
  in
    Array.fast_sort (compare) a;
    contains_array := a;
    Array.iter (fun x -> Hashtbl.add contains_index x true) !contains_array
      
let count_contains t s = text_count_contains t.doc s
let unsorted_contains t s = text_unsorted_contains t.doc s

let init_naive_contains t s =
  let i,j = tree_doc_ids t.doc (tree_root t.doc)
  in
  let regexp = Str.regexp_string s in
  let matching arg = 
    try
      let _ = Str.search_forward regexp arg 0;
      in true
    with _ -> false
  in
  let rec loop n acc l = 
    if n >= j then acc,l
    else
      let s = text_get_cached_text t.doc n
      in
	if matching s 
	then loop (nodei ((inode n)+1)) (n::acc) (l+1) 
	else loop (nodei ((inode n)+1)) acc l
  in
  let acc,l = loop i [] 0 in
  let a = Array.create l nulldoc in
  let _ = List.fold_left (fun cpt e -> a.(cpt) <- e; (cpt-1)) (l-1) acc
  in
    contains_array := a

let last_idx = ref 0

let array_find a i j =
  let l = Array.length a in
  let rec loop idx x y =
    if x > y || idx >= l then nulldoc
       else
         if a.(idx) >= x then if a.(idx) > y then nulldoc else (last_idx := idx;a.(idx))
         else loop (idx+1) x y
  in
    if a.(0) > j || a.(l-1) < i then nulldoc
    else loop !last_idx i j 
	  
let text_below tree t = 
  let l = Array.length !contains_array in
  let i,j = tree_doc_ids tree.doc t in
  let id = if l == 0 then i else (array_find !contains_array i j) in
  tree_parent_doc tree.doc id
    
let text_next tree t root =
  let l = Array.length !contains_array in
  let inf = nodei((inode(snd(tree_doc_ids tree.doc t)))+1) in    
  let _,j = tree_doc_ids tree.doc root in      
  let id = if l == 0 then if inf > j then nulldoc else  inf
  else array_find !contains_array inf j
  in 
  tree_parent_doc tree.doc id



module DocIdSet = struct
  include Set.Make (struct type t = [`Text] node
			   let compare = compare_node end)
    
end
let is_nil t = t == nil

let is_node t = t != nil
let is_root t = t == root

let node_of_t t  =
  let _ = Tag.init (Obj.magic t) in
  let table = collect_tags t 
  in (*
  let _ = Hashtbl.iter (fun t (c,d,ns,f) ->
			  Printf.eprintf "Tag %s has:\n" (Tag.to_string t);
			  Printf.eprintf "Child tags: ";
			  Ptset.Int.iter (fun t -> Printf.eprintf "%s "(Tag.to_string t)) c;
			  Printf.eprintf "\nDescendant tags: ";
			  Ptset.Int.iter (fun t -> Printf.eprintf "%s "(Tag.to_string t)) d;
			  Printf.eprintf "\nNextSibling tags: ";
			  Ptset.Int.iter (fun t -> Printf.eprintf "%s "(Tag.to_string t)) ns;
			  Printf.eprintf "\nFollowing tags: ";
			  Ptset.Int.iter (fun t -> Printf.eprintf "%s "(Tag.to_string t)) f;
			  Printf.eprintf "\n\n%!";) table
  in
			  
     *)			  
    { doc= t; 
      ttable = table;
    }

let finalize _ = Printf.eprintf "Release the string list !\n%!"
;;

let parse f str =
  node_of_t
    (f str 
       !Options.sample_factor 
       !Options.index_empty_texts
       !Options.disable_text_collection)
    
let parse_xml_uri str = parse parse_xml_uri str
let parse_xml_string str =  parse parse_xml_string str

     
external pool : tree -> Tag.pool = "%identity"

let magic_string = "SXSI_INDEX"
let version_string = "2"

let pos fd =
  Unix.lseek fd 0  Unix.SEEK_CUR

let pr_pos fd = Printf.eprintf "At position %i\n%!" (pos fd)

let write fd s = 
  let sl = String.length s in
  let ssl = Printf.sprintf "%020i" sl in
    ignore (Unix.write fd ssl 0 20);
    ignore (Unix.write fd s 0 (String.length s))

let rec really_read fd buffer start length =
  if length <= 0 then () else
    match Unix.read fd buffer start length with
	0 -> raise End_of_file
      | r -> really_read fd buffer (start + r) (length - r);;

let read fd =
  let buffer = String.create 20 in
  let _ =  really_read fd buffer 0 20 in
  let size = int_of_string buffer in
  let buffer = String.create size in
  let _ =  really_read fd buffer 0 size in
    buffer
    

let save t str =
  let fd = Unix.openfile str [ Unix.O_WRONLY;Unix.O_TRUNC;Unix.O_CREAT] 0o644 in
  let out_c = Unix.out_channel_of_descr fd in
  let _ = set_binary_mode_out out_c true in
    output_string out_c magic_string;
    output_char out_c '\n';
    output_string out_c version_string;
    output_char out_c '\n';
    Marshal.to_channel out_c t.ttable [ ];
    (* we need to move the fd to the correct position *)
    flush out_c;
    ignore (Unix.lseek fd (pos_out out_c) Unix.SEEK_SET);
    tree_save t.doc fd;
    close_out out_c
;;

let load ?(sample=64) str = 
  let fd = Unix.openfile str [ Unix.O_RDONLY ] 0o644 in
  let in_c = Unix.in_channel_of_descr fd in
  let _ = set_binary_mode_in in_c true in
  let load_table () = 
    (let ms = input_line in_c in if ms <> magic_string then failwith "Invalid index file");
    (let vs = input_line in_c in if vs <> version_string then failwith "Invalid version file");
    let table : (Tag.t,(Ptset.Int.t*Ptset.Int.t*Ptset.Int.t*Ptset.Int.t)) Hashtbl.t =
      Marshal.from_channel in_c 
    in
    let ntable = Hashtbl.create (Hashtbl.length table) in
      Hashtbl.iter (fun k (s1,s2,s3,s4) -> 
		      let ss1 = Ptset.Int.fold (Ptset.Int.add) s1 Ptset.Int.empty
		      and ss2 = Ptset.Int.fold (Ptset.Int.add) s2 Ptset.Int.empty
		      and ss3 = Ptset.Int.fold (Ptset.Int.add) s3 Ptset.Int.empty
		      and ss4 = Ptset.Int.fold (Ptset.Int.add) s4 Ptset.Int.empty
		      in Hashtbl.add ntable k (ss1,ss2,ss3,ss4)
		   ) table;
      Hashtbl.clear table;
      (* The in_channel read a chunk of fd, so we might be after
	 the start of the XMLTree save file. Reset to the correct
	 position *)
      ntable
  in
  let _ = Printf.eprintf "\nLoading tag table : " in
  let ntable = time (load_table) () in
  ignore(Unix.lseek fd (pos_in in_c) Unix.SEEK_SET);
  let tree = { doc = tree_load fd;
	       ttable = ntable;}
  in close_in in_c;
  tree
  



let tag_pool t = pool t.doc
  
let compare = compare_node

let equal a b = a == b
   
let nts = function
    -1 -> "Nil"
  | i -> Printf.sprintf "Node (%i)"  i
      
let dump_node t = nts (inode t)

let is_left t n = tree_is_first_child t.doc n

let is_below_right t n1 n2 = 
  tree_is_ancestor t.doc (tree_parent t.doc n1) n2 
  && not (tree_is_ancestor t.doc n1 n2)
    
let parent t n = tree_parent t.doc n

let first_child t = (); fun n -> tree_first_child t.doc n
let first_element t = (); fun n -> tree_first_element t.doc n

(* these function will be called in two times: first partial application
   on the tag, then application of the tag and the tree, then application of
   the other arguments. We use the trick to let the compiler optimize application
*)

let tagged_child t tag = () ; fun n -> tree_tagged_child t.doc n tag

let select_child t = fun ts ->
  let v = ptset_to_vector ts in ();
    fun n -> tree_select_child t.doc n v

let next_sibling t = (); fun n ->  tree_next_sibling t.doc n
let next_element t = (); fun n ->  tree_next_element t.doc n

let tagged_sibling t tag = (); fun n -> tree_tagged_sibling t.doc n tag

let select_sibling t = fun ts ->
  let v = (ptset_to_vector ts) in ();
    fun n -> tree_select_foll_sibling t.doc n v

let next_sibling_ctx t = (); fun n _ -> tree_next_sibling t.doc n
let next_element_ctx t = (); fun n _ ->  tree_next_element t.doc n
let tagged_sibling_ctx t tag = (); fun n  _ -> tree_tagged_sibling t.doc n tag

let select_sibling_ctx t = fun ts -> 
  let v = (ptset_to_vector ts) in ();
     fun n  _ -> tree_select_foll_sibling t.doc n v

let id t n = tree_node_xml_id t.doc n
	
let tag t n = if n == nil then Tag.nullt else tree_tag_id t.doc n

let tagged_desc t tag = (); fun n -> tree_tagged_desc t.doc n tag 

let select_desc t = fun ts -> 
  let v = (ptset_to_vector ts) in ();
    fun n -> tree_select_desc t.doc n v

let tagged_foll_ctx  t tag = (); fun n ctx -> tree_tagged_foll_below t.doc n tag ctx

let select_foll_ctx t = fun ts ->
  let v = (ptset_to_vector ts) in ();
    fun n ctx -> tree_select_foll_below t.doc n v ctx

let last_idx = ref 0
let array_find a i j =
  let l = Array.length a in
  let rec loop idx x y =
    if x > y || idx >= l then nil
	else
	  if a.(idx) >= x then if a.(idx) > y then nil else (last_idx := idx;a.(idx))
	  else loop (idx+1) x y
  in
    if a.(0) > j || a.(l-1) < i then nil
    else loop !last_idx i j 



  let count t s = text_count t.doc s

  let print_xml_fast outc tree t =
    let rec loop ?(print_right=true) t = 
      if t != nil 
      then 
	let tagid = tree_tag_id tree.doc t in
	  if tagid==Tag.pcdata
	  then 
	    begin
 	      let tid =  tree_my_text tree.doc t in
	      output_string outc (text_get_cached_text tree.doc tid);
	      if print_right
	      then loop (next_sibling tree t);
	    end
	  else
	    let tagstr = Tag.to_string tagid in
	    let l = first_child tree t 
	    and r = next_sibling tree t 
	    in
	      output_char outc  '<';
	      output_string outc  tagstr;
	      if l == nil then output_string outc  "/>"
	      else 
		if (tag tree l) == Tag.attribute then
		  begin
		    loop_attributes (first_child tree l);
		    if (next_sibling tree l) == nil then output_string outc  "/>"
		    else  
		      begin 
			output_char outc  '>'; 
			loop (next_sibling tree l);
			output_string outc  "</";
			output_string outc  tagstr;
			output_char outc '>';
		      end;
		  end
		else
		  begin
		    output_char outc  '>'; 
		    loop l;
		    output_string outc "</";
		    output_string outc tagstr;
		    output_char outc '>';
		  end;
	      if print_right then loop r
    and loop_attributes a = 
      if a != nil
      then
      let s = (Tag.to_string (tag tree a)) in
      let attname = String.sub s 3 ((String.length s) -3) in
      let fsa = first_child tree a in
      let tid =  tree_my_text tree.doc fsa in
	output_char outc ' ';
	output_string outc attname;
	output_string outc "=\"";
	output_string outc (text_get_cached_text tree.doc tid);
	output_char outc '"';
	loop_attributes (next_sibling tree a)
    in
	loop ~print_right:false t
	  
	  
    let print_xml_fast outc tree t = 
      if (tag tree t) = Tag.document_node then
	print_xml_fast outc tree (first_child tree t)
      else print_xml_fast outc tree t 
	
let tags_children t tag = 
  let a,_,_,_ = Hashtbl.find t.ttable tag in a
let tags_below t tag = 
  let _,a,_,_ = Hashtbl.find t.ttable tag in a
let tags_siblings t tag = 
  let _,_,a,_ = Hashtbl.find t.ttable tag in a
let tags_after t tag = 
  let _,_,_,a = Hashtbl.find t.ttable tag in a


let tags t tag = Hashtbl.find t.ttable tag


let rec binary_parent t n = 
  let r = 
  if tree_is_first_child t.doc n
  then tree_parent t.doc n
  else tree_prev_sibling t.doc n
  in if tree_tag_id t.doc r = Tag.pcdata then
  binary_parent t r
  else r

let doc_ids t n = tree_doc_ids t.doc n

let subtree_tags t tag = ();
  fun n -> if n == nil then 0 else
    tree_subtree_tags t.doc n tag

let get_text t n =
  let tid = tree_my_text t.doc n in
    if tid == nulldoc then "" else 
      text_get_cached_text t.doc tid


let dump_tree fmt tree = 
  let rec loop t n =
    if t != nil then
      let tag = (tree_tag_id tree.doc t ) in
      let tagstr = Tag.to_string tag in
	let tab = String.make n ' ' in

	  if tag == Tag.pcdata || tag == Tag.attribute_data 
	  then 
	    Format.fprintf fmt "%s<%s>%s</%s>\n" 
	      tab tagstr (text_get_cached_text tree.doc (tree_my_text tree.doc t)) tagstr
	  else begin
	    Format.fprintf fmt "%s<%s>\n" tab tagstr;
	    loop (tree_first_child tree.doc t) (n+2);
	    Format.fprintf fmt "%s</%s>\n%!" tab tagstr;
	  end;
	  loop (tree_next_sibling tree.doc t) n
  in
    loop root 0
;;