(******************************************************************************)
(*  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"

type tree
type 'a node = int
type node_kind = [`Text | `Tree ]
    
let compare_node : 'a node -> 'a node -> int = (-)
let equal_node : 'a node -> 'a node -> bool = (==)
  
(* abstract type, values are pointers to a XMLTree C++ object *)

external int_of_node : 'a node -> int = "%identity"
  
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 save_tree : tree -> string -> unit = "caml_xml_tree_save"
external load_tree : string ->  int -> tree = "caml_xml_tree_load"
  
external nullt : unit -> 'a node = "caml_xml_tree_nullt"

let nil : 'a node = Obj.magic (-1)

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 nil 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_serialize : tree -> string -> unit = "caml_xml_tree_serialize"

external tree_unserialize : string -> tree = "caml_xml_tree_unserialize"
      
external tree_root : tree -> [`Tree] node = "caml_xml_tree_root" 

let tree_is_nil x = equal_node x nil

external tree_parent : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_parent" 
external tree_parent_doc : tree -> [`Text ] node -> [`Tree ] node = "caml_xml_tree_parent_doc" 
external tree_prev_doc : tree -> [`Text ] node -> [`Tree ] node = "caml_xml_tree_prev_doc" 
external tree_first_child : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_first_child" 
external tree_tagged_child : tree -> [`Tree] node -> Tag.t -> [`Tree] node = "caml_xml_tree_tagged_child" 
external tree_next_sibling : tree -> [`Tree] node -> [`Tree] node = "caml_xml_tree_next_sibling" 
external tree_tagged_sibling : tree -> [`Tree] node -> Tag.t -> [`Tree] node = "caml_xml_tree_tagged_sibling" 

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

(*    external tag : tree -> [`Tree ] node -> T = "caml_xml_tree_tag"*)
external tree_tag_id : tree -> [`Tree ] node -> Tag.t = "caml_xml_tree_tag_id" 
    

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" 

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

let text_size tree = int_of_node (snd ( tree_doc_ids tree (Obj.magic 0) ))

let text_get_cached_text t x =
  if x == -1 then ""
  else 
     text_get_cached_text t x


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



type int_vector
external int_vector_alloc : int -> int_vector = "caml_int_vector_alloc"
external int_vector_length : int_vector -> int = "caml_int_vector_length"
external int_vector_set : int_vector -> int -> int -> unit = "caml_int_vector_set"

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


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 = int_vector_alloc (Ptset.Int.cardinal s) in
	let _ = Ptset.Int.fold (fun e i -> int_vector_set v i e;i+1) s 0 in
	  HPtset.add vector_htbl s v; v

      
type t = { doc : tree;	          
	   node : [`Tree] node;
	   ttable : (Tag.t,(Ptset.Int.t*Ptset.Int.t)) Hashtbl.t;
	 }

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)

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 = Hashtbl.create BIG_H_SIZE in
  let pt_add t s = 
    let k = HASHINT2(Tag.hash t,Ptset.Int.hash s) in
      try
	Hashtbl.find h_add k
      with
      | Not_found -> let r = Ptset.Int.add t s in
	  Hashtbl.add h_add k r;r
  in
  let h = Hashtbl.create BIG_H_SIZE in
  let update t sb sa =
    let sbelow,safter = 
      try
	Hashtbl.find h t 
      with
	| Not_found -> 
	    (Ptset.Int.empty,Ptset.Int.empty)
    in
      Hashtbl.replace h t (pt_cup sbelow sb, pt_cup safter sa)
  in
  let rec loop id acc = 
    if equal_node id nil
    then (Ptset.Int.empty,acc)
    else
      let below2,after2 = loop (tree_next_sibling tree id) acc in
      let below1,after1 = loop (tree_first_child tree id) after2 in
      let tag = tree_tag_id tree id in
	update tag below1 after2;
	pt_add tag (pt_cup below1 below2), (pt_add tag after1)
  in
    let _ = loop (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 (n+1) (n::acc) (l+1) 
	else loop (n+1) acc l
  in
  let acc,l = loop i [] 0 in
  let a = Array.create l nil in
  let _ = List.fold_left (fun cpt e -> a.(cpt) <- e; (cpt-1)) (l-1) acc
  in
    contains_array := a
	  


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

let is_node t = t.node != nil

let node_of_t t  =
  let _ = Tag.init (Obj.magic t) in
  let table = collect_tags t 
  in
    { doc= t; 
      node = tree_root 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 save t str = (save_tree t.doc str)
;;

let load ?(sample=64) str = 
  node_of_t (load_tree str sample)
    



let tag_pool t = pool t.doc
  
let compare a b = a.node - b.node

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

let mk_nil t = { t with node = nil }		  
let root n = { n with node = tree_root n.doc }

let is_root n = n.node == (tree_root n.doc)
      
let is_left n = tree_is_first_child n.doc n.node

let is_below_right t1 t2 = tree_is_ancestor t1.doc (tree_parent t1.doc t1.node) t2.node

let parent n =  { n with node = tree_parent n.doc n.node }

let first_child n = { n with node = tree_first_child n.doc n.node }
let tagged_child tag n  =  { n with node = tree_tagged_child n.doc n.node tag }
let select_child ts n  =  { n with node = tree_select_child n.doc n.node (ptset_to_vector ts) }

let next_sibling n = { n with node = tree_next_sibling n.doc n.node }
let tagged_sibling tag n  =  { n with node = tree_tagged_sibling n.doc n.node tag }
let select_sibling ts n  =  { n with node = tree_select_foll_sibling n.doc n.node (ptset_to_vector ts) }

let next_sibling_ctx n _ = next_sibling n
let tagged_sibling_ctx tag n  _ = tagged_sibling tag n
let select_sibling_ctx ts n  _ = select_sibling ts n

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

let tagged_desc tag n = { n with node = tree_tagged_desc n.doc n.node tag }
let select_desc ts n  =  { n with node = tree_select_desc n.doc n.node (ptset_to_vector ts) }

let tagged_foll_ctx tag t ctx =
  { t with node = tree_tagged_foll_below t.doc t.node tag ctx.node }
let select_foll_ctx ts n ctx  =  { n with node = tree_select_foll_below n.doc n.node (ptset_to_vector ts) ctx.node }

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 t =
    let rec loop ?(print_right=true) t = 
      if t.node != nil 
      then 
	let tagid = tree_tag_id t.doc t.node in
	  if tagid==Tag.pcdata
	  then output_string outc (text_get_cached_text t.doc t.node);
	  if print_right
	  then loop (next_sibling t)
	    
	  else
	    let tagstr = Tag.to_string tagid in
	    let l = first_child t 
	    and r = next_sibling t 
	    in
	      output_char outc  '<';
	      output_string outc  tagstr;
	      if l.node == nil then output_string outc  "/>"
	      else 
		if (tag l) == Tag.attribute then
		  begin
		    loop_attributes (first_child l);
		    if (next_sibling l).node == nil then output_string outc  "/>"
		    else  
		      begin 
			output_char outc  '>'; 
			loop (next_sibling 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 =    
      let s = (Tag.to_string (tag a)) in
      let attname = String.sub s 3 ((String.length s) -3) in
	output_char outc ' ';
	output_string outc attname;
	output_string outc "=\"";
	output_string outc (text_get_cached_text t.doc
			      (tree_my_text a.doc (first_child a).node));
	output_char outc '"';
	loop_attributes (next_sibling a)
    in
	loop ~print_right:false t
	  
	  
    let print_xml_fast outc t = 
      if (tag t) = Tag.document_node then
	print_xml_fast outc (first_child t)
      else print_xml_fast outc t 
	

let tags_below t tag = 
  fst(Hashtbl.find t.ttable tag)

let tags_after t tag = 
  snd(Hashtbl.find t.ttable tag)

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


let rec binary_parent t = 
  if tree_is_first_child t.doc t.node
  then { t with node = tree_parent t.doc t.node }
  else { t with node = tree_prev_sibling t.doc t.node }

let doc_ids (t:t) : (int*int) = 
  (Obj.magic (tree_doc_ids t.doc t.node))

let subtree_tags t tag = 
  if t.node == nil then 0 else
    tree_subtree_tags t.doc t.node tag

let get_text t =
  let tid = tree_my_text t.doc t.node in
    if tid == nil then "" else 
      let a, b = tree_doc_ids t.doc (tree_root t.doc) in
      let _ = Printf.eprintf "Trying to take text %i of node %i in %i %i\n%!" tid t.node a b in
	text_get_cached_text t.doc tid


let dump_tree fmt t = 
  let rec loop tree n =
    if tree != nil then
      let tag = (tree_tag_id t.doc tree ) 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 t.doc (tree_my_text t.doc tree)) tagstr
	  else begin
	    Format.fprintf fmt "%s<%s>\n" tab tagstr;
	    loop (tree_first_child t.doc tree) (n+2);
	    Format.fprintf fmt "%s</%s>\n%!" tab tagstr;
	  end;
	  loop (tree_next_sibling t.doc tree) n
  in
    loop (tree_root t.doc) 0
;;