af2ec3f259596dac27266d988b5eac22f638e4a8
[SXSI/xpathcomp.git] / src / tree.ml
1 (******************************************************************************)
2 (*  SXSI : XPath evaluator                                                    *)
3 (*  Kim Nguyen (Kim.Nguyen@nicta.com.au)                                      *)
4 (*  Copyright NICTA 2008                                                      *)
5 (*  Distributed under the terms of the LGPL (see LICENCE)                     *)
6 (******************************************************************************)
7 INCLUDE "debug.ml"
8   INCLUDE "utils.ml"
9
10 external init_lib : unit -> unit = "sxsi_cpp_init"
11
12 exception CPlusPlusError of string
13
14 let () = Callback.register_exception "CPlusPlusError" (CPlusPlusError "")
15
16 let () =  init_lib ()
17
18 type node = [ `Tree ] Node.t
19
20 type tree
21
22 external register_tag : tree -> string -> Tag.t = "caml_xml_tree_register_tag"
23
24 external tag_name : tree -> Tag.t -> string = "caml_xml_tree_get_tag_name"
25
26 let tag t = (); fun s ->
27   match s with
28   | "<$>" -> Tag.pcdata
29   | "<@>" -> Tag.attribute
30   | "" -> Tag.document_node
31   | "<@$>" -> Tag.attribute_data
32   | _ -> register_tag t s
33
34 let to_string d = ();
35   fun t ->
36     if t == Tag.pcdata then "<$>"
37     else if t == Tag.attribute_data then "<@$>"
38     else if t == Tag.attribute then "<@>"
39     else if t == Tag.nullt then "<!NIL!>"
40     else tag_name d t
41
42 let translate  x = x
43
44 let mk_tag_ops t = {
45   Tag.tag = tag t;
46   Tag.to_string = to_string t;
47   Tag.translate = translate
48 }
49
50 module TreeBuilder =
51 struct
52   type t
53   external create : unit -> t = "caml_xml_tree_builder_create"
54   external open_document : t -> int -> bool -> int -> unit = "caml_xml_tree_builder_open_document"
55   external close_document : t -> tree = "caml_xml_tree_builder_close_document"
56   external open_tag : t -> string -> unit = "caml_xml_tree_builder_open_tag"
57   external close_tag : t -> string -> unit = "caml_xml_tree_builder_close_tag"
58   external text : t -> string -> unit = "caml_xml_tree_builder_text"
59
60   let is_whitespace s =
61     let rec loop len i =
62       if i < len then
63         let c = s.[i] in
64         (c == '\n' || c == '\t' || c == ' ') && loop len (i+1)
65       else
66         true
67     in
68     loop (String.length s) 0
69
70
71   let display_count =
72     let event_counter = ref 0 in
73     (fun parser_ ->
74       incr event_counter;
75       if !event_counter land 0xffffff == 0 then
76         Logger.print Format.err_formatter "Current position: %i@\n@?" (Expat.get_current_byte_index parser_))
77
78
79   let do_text b t =
80     if Buffer.length t > 0 then begin
81       let s = Buffer.contents t in
82       begin
83         open_tag b "<$>";
84         text b s;
85         close_tag b "<$>";
86       end;
87       Buffer.clear t
88     end
89
90   let output_attr b name value =
91     let atname = "<@>" ^ name in
92     open_tag b atname;
93     open_tag b "<@$>";
94     text b value;
95     close_tag b "<@$>";
96     close_tag b atname
97
98   let start_element_handler parser_ b t tag attr_list =
99     do_text b t;
100     open_tag b tag;
101     match attr_list with
102       [] -> ()
103     | l ->
104       open_tag b "<@>";
105       List.iter (fun (name, value) -> output_attr b name value) l;
106       close_tag b "<@>"
107
108
109   let end_element_handler parser_ b t tag =
110     do_text b t;
111     close_tag b tag
112
113   let character_data_handler parser_ _ t text =
114     Buffer.add_string t text
115
116   let create_parser () =
117     let buf = Buffer.create 512 in
118     let build = create () in
119     let parser_ = Expat.parser_create ~encoding:None in
120     let finalize () =
121       do_text build buf;
122       close_tag build "";
123       Logger.print Format.err_formatter "Finished parsing@\n";
124       Logger.print Format.err_formatter "Starting index construction@\n";
125       let r = close_document build in
126       Logger.print Format.err_formatter "Finished index construction@\n";
127       r
128     in
129     Expat.set_start_element_handler parser_ (start_element_handler parser_ build buf);
130     Expat.set_end_element_handler parser_ (end_element_handler parser_ build buf);
131     Expat.set_character_data_handler parser_ (character_data_handler parser_ build buf);
132     Logger.print Format.err_formatter "Started parsing@\n";
133     open_document build !Options.sample_factor !Options.disable_text_collection !Options.text_index_type;
134     open_tag build "";
135     parser_, finalize
136
137   let parse_string s =
138     let parser_, finalizer = create_parser () in
139     Expat.parse parser_ s;
140     finalizer ()
141
142   let parse_file file =
143     let in_chan = open_in file in
144     let buffer = String.create 4096 in
145     let parser_, finalizer = create_parser () in
146     let () =
147       try
148         while true do
149           let read = input in_chan buffer 0 4096 in
150           if read == 0 then raise End_of_file else
151             Expat.parse_sub parser_ buffer 0 read;
152         done
153
154       with
155       | End_of_file -> close_in in_chan
156       | e -> raise e
157     in
158     finalizer ()
159
160 end
161
162
163
164
165 type bit_vector = string
166
167 external bool_of_int : int -> bool = "%identity"
168 external int_of_bool : bool -> int = "%identity"
169
170 let bit_vector_unsafe_get v i =
171   bool_of_int
172     (((Char.code (String.unsafe_get v (i lsr 3))) lsr (i land 7)) land 1)
173 let chr (c:int) : char = Obj.magic (c land 0xff)
174 let bit_vector_unsafe_set v i b =
175   let j = i lsr 3 in
176   let c = Char.code v.[j] in
177   let bit = int_of_bool b in
178   let mask = bit lsl (i land 7) in
179   if b then v.[j] <- chr (c lor mask) else v.[j] <- (chr (c land (lnot mask)))
180
181 let bit_vector_create n =
182   let len = if n <= 0 then 0 else (n - 1) / 8 + 1 in
183   String.make len '\000'
184
185 type t = {
186   doc : tree;
187   elements: Ptset.Int.t;
188   attributes: Ptset.Int.t;
189   attribute_array : Tag.t array;
190   children : Ptset.Int.t array;
191   siblings : Ptset.Int.t array;
192   descendants: Ptset.Int.t array;
193   followings: Ptset.Int.t array;
194 }
195
196
197 let tag_operations t = mk_tag_ops t.doc
198 (*
199   external parse_xml_uri : string -> int -> bool -> bool -> int -> tree = "caml_call_shredder_uri"
200   external parse_xml_string :  string -> int -> bool -> bool -> int -> tree = "caml_call_shredder_string"
201 *)
202 external tree_print_xml_fast3 : tree -> [`Tree ] Node.t -> Unix.file_descr -> unit = "caml_xml_tree_print"
203 let print_xml t n fd =
204   tree_print_xml_fast3 t.doc n fd
205
206
207 external tree_save : tree -> Unix.file_descr -> string -> unit = "caml_xml_tree_save"
208 external tree_load : Unix.file_descr -> string -> bool -> int -> tree = "caml_xml_tree_load"
209
210 external nullt : unit -> 'a Node.t = "caml_xml_tree_nullt"
211
212 let nil : [`Tree ] Node.t = Node.nil
213 let root : [`Tree ] Node.t = Node.null
214
215 type tag_list
216
217 external tag_list_alloc : int -> tag_list = "caml_tag_list_alloc"
218 external tag_list_set : tag_list -> int -> Tag.t -> unit = "caml_tag_list_set" "noalloc"
219
220 module HPtset = Hashtbl.Make(Ptset.Int)
221
222 let vector_htbl = HPtset.create MED_H_SIZE
223 let reinit () = HPtset.clear vector_htbl
224
225 let tag_list_of_set s =
226   try
227     HPtset.find vector_htbl s
228   with
229     Not_found ->
230       let v = tag_list_alloc (Ptset.Int.cardinal s + 1) in
231       let i = ref 0 in
232       let () = Ptset.Int.iter (fun e -> tag_list_set v !i e; incr i) s in
233       let () = tag_list_set v !i Tag.nullt in
234       HPtset.add vector_htbl s v; v
235
236 (** tree interface *)
237
238 external tree_root : tree -> [`Tree] Node.t = "caml_xml_tree_root"  "noalloc"
239
240
241 external tree_first_child : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_first_child" "noalloc"
242 let first_child t n = tree_first_child t.doc n
243
244 external tree_first_element : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_first_element" "noalloc"
245 let first_element t n = tree_first_element t.doc n
246
247 external tree_tagged_child : tree -> [`Tree] Node.t -> Tag.t -> [`Tree] Node.t = "caml_xml_tree_tagged_child" "noalloc"
248 let tagged_child t n tag = tree_tagged_child t.doc n tag
249
250 external tree_select_child : tree -> [`Tree ] Node.t -> tag_list -> [`Tree] Node.t = "caml_xml_tree_select_child" "noalloc"
251 let select_child t n tag_set = tree_select_child t.doc n tag_set
252
253 external tree_last_child : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_last_child" "noalloc"
254 let last_child t n = tree_last_child t.doc n
255
256
257 external tree_next_sibling : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_next_sibling"  "noalloc"
258 let next_sibling t n = tree_next_sibling t.doc n
259
260 external tree_next_element : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_next_element"  "noalloc"
261 let next_element t n = tree_next_element t.doc n
262
263
264 external tree_tagged_sibling : tree -> [`Tree] Node.t -> Tag.t -> [`Tree] Node.t = "caml_xml_tree_tagged_sibling" "noalloc"
265 let tagged_sibling t n tag = tree_tagged_sibling t.doc n tag
266
267
268 external tree_select_sibling : tree -> [`Tree ] Node.t -> tag_list -> [`Tree] Node.t = "caml_xml_tree_select_sibling" "noalloc"
269 let select_sibling t n tag_set = tree_select_sibling t.doc n tag_set
270
271 external tree_prev_sibling : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_prev_sibling" "noalloc"
272 let prev_sibling t n = tree_prev_sibling t.doc n
273
274
275
276 external tree_tagged_descendant : tree -> [`Tree ] Node.t -> Tag.t -> [`Tree ] Node.t = "caml_xml_tree_tagged_descendant" "noalloc"
277 let tagged_descendant t n tag = tree_tagged_descendant t.doc n tag
278
279 external tree_tagged_next : tree -> [`Tree ] Node.t -> Tag.t -> [`Tree ] Node.t = "caml_xml_tree_tagged_next" "noalloc"
280 let tagged_next t n tag = tree_tagged_next t.doc n tag
281
282 external tree_select_descendant : tree -> [`Tree ] Node.t -> tag_list -> [`Tree] Node.t = "caml_xml_tree_select_descendant" "noalloc"
283 let select_descendant t n tag_set = tree_select_descendant t.doc n tag_set
284
285 external tree_tagged_following_before : tree -> [`Tree ] Node.t -> Tag.t -> [`Tree ] Node.t -> [`Tree ] Node.t = "caml_xml_tree_tagged_following_before" "noalloc"
286 let tagged_following_before t n tag ctx = tree_tagged_following_before t.doc n tag ctx
287
288 external tree_select_following_before : tree -> [`Tree ] Node.t -> tag_list -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_select_following_before" "noalloc"
289 let select_following_before t n tag_set ctx = tree_select_following_before t.doc n tag_set ctx
290
291 external tree_parent : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_parent" "noalloc"
292 let parent t n = tree_parent t.doc n
293
294 external tree_tag : tree -> [`Tree] Node.t -> Tag.t = "caml_xml_tree_tag" "noalloc"
295 let tag t n = tree_tag t.doc n
296
297 external tree_is_first_child : tree -> [ `Tree ] Node.t -> bool = "caml_xml_tree_is_first_child" "noalloc"
298 let is_first_child t n = tree_is_first_child t.doc n
299
300 external tree_is_right_descendant : tree -> [ `Tree ] Node.t -> [`Tree] Node.t -> bool =
301   "caml_xml_tree_is_right_descendant" "noalloc"
302
303 let is_right_descendant t n1 n2 = tree_is_right_descendant t.doc n1 n2
304 ;;
305
306 let node_tags t = Ptset.Int.add Tag.document_node t.descendants.(Tag.document_node)
307
308 let attribute_tags t = t.attributes
309
310 let element_tags t = t.elements
311
312 let tags t tag =
313   t.children.(tag), t.descendants.(tag), t.siblings.(tag), t.followings.(tag)
314
315 open Format
316 let dump_tag_table t =
317   let tag = ref 0 in
318   let printer ppf set =
319     Logger.print ppf "%s: %a"
320       (Tag.to_string !tag) TagSet.print (TagSet.inj_positive set);
321     incr tag
322   in
323   let set_printer msg set =
324     tag := 0;
325     Logger.print err_formatter "%s :@\n" msg;
326     Pretty.pp_print_array ~sep:pp_force_newline printer err_formatter set;
327     Logger.print err_formatter "-----------------------------@\n";
328   in
329   set_printer "Child tags" t.children;
330   set_printer "Descendant tags" t.descendants;
331   set_printer "Sibling tags" t.siblings;
332   set_printer "Following tags" t.followings
333
334 external tree_subtree_tags : tree -> [`Tree] Node.t -> Tag.t -> int = "caml_xml_tree_subtree_tags" "noalloc"
335 let subtree_tags t n tag = tree_subtree_tags t.doc n tag
336
337 external tree_subtree_size : tree -> [`Tree] Node.t -> int = "caml_xml_tree_subtree_size" "noalloc"
338 let subtree_size t n = tree_subtree_size t.doc n
339
340 external tree_subtree_elements : tree -> [`Tree] Node.t -> int = "caml_xml_tree_subtree_elements" "noalloc"
341 let subtree_elements t n = tree_subtree_elements t.doc n
342
343 external tree_closing : tree -> [`Tree] Node.t -> [`Tree] Node.t = "caml_xml_tree_closing" "noalloc"
344 let closing t n = tree_closing t.doc n
345
346 external tree_num_tags : tree -> int = "caml_xml_tree_num_tags" "noalloc"
347 let num_tags t = tree_num_tags t.doc
348
349 external tree_size : tree -> int = "caml_xml_tree_size" "noalloc"
350 let size t = tree_size t.doc
351
352
353 let stats t =
354   let tree = t.doc in
355   let rec loop left node acc_d total_d num_leaves =
356     if node == nil then
357       (acc_d+total_d,if left then num_leaves+1 else num_leaves)
358     else
359       let d,td = loop true (tree_first_child tree node) (acc_d+1) total_d num_leaves in
360       loop false (tree_next_sibling tree  node) (acc_d)  d td
361   in
362   let a,b = loop true root 0 0 0
363   in
364   Logger.print err_formatter "Average depth: %f, number of leaves %i@\n@?" ((float_of_int a)/. (float_of_int b)) b
365 ;;
366
367 module TagS =
368 struct
369   include Ptset.Make (
370     struct type t = int
371            type data = t
372            external hash : t -> int = "%identity"
373            external uid : t -> Uid.t = "%identity"
374            external equal : t -> t -> bool = "%eq"
375            external make : t -> int = "%identity"
376            external node : t -> int = "%identity"
377            external stats : unit -> unit = "%identity"
378            external init : unit -> unit = "%identity"
379     end
380   )
381   let to_ptset s = fold (Ptset.Int.add) s Ptset.Int.empty
382 end
383
384 module TSTSCache =
385   Hashtbl.Make(struct type t = TagS.t * TagS.t
386                       let hash (x, y) =
387                         HASHINT2(Uid.to_int x.TagS.Node.id,
388                                  Uid.to_int y.TagS.Node.id)
389                       let equal u v =
390                         let u1,u2 = u
391                         and v1,v2 = v in
392                         u1 == v1 && u2 == v2
393   end)
394 module TagTSCache =
395   Hashtbl.Make(struct type t = Tag.t * TagS.t
396                       let hash (x, y) =
397                         HASHINT2(x, Uid.to_int y.TagS.Node.id)
398                       let equal u v =
399                         let u1,u2 = u
400                         and v1,v2 = v in
401                         u1 == v1 && u2 == v2
402   end)
403
404 let add_cache = TagTSCache.create 1023
405 let union_cache = TSTSCache.create 1023
406 let subset_cache = TSTSCache.create 1023
407
408 let clear_cache () =
409   TSTSCache.clear union_cache;
410   TSTSCache.clear subset_cache;
411   TagTSCache.clear add_cache
412
413 let _subset x y =
414   (x == y) || (x == TagS.empty) ||
415     if y == TagS.empty then false
416     else
417       let key = (x, y) in
418       try
419         TSTSCache.find subset_cache key
420       with
421       | Not_found ->
422         let z = TagS.subset x y in
423         TSTSCache.add subset_cache key z; z
424
425 let order ((x, y) as z) =
426   if x.TagS.Node.id <= y.TagS.Node.id then z
427   else (y, x)
428
429 let _union x y =
430   if _subset x y then y
431   else if _subset y x then x
432   else
433     let key = order (x, y) in
434     try
435       TSTSCache.find union_cache key
436     with
437     | Not_found ->
438       let z = TagS.union x y in
439       TSTSCache.add union_cache key z; z
440
441 let _add t s =
442   let key = (t,s) in
443   try
444     TagTSCache.find add_cache key
445   with
446   | Not_found ->
447     let z = TagS.add t s in
448     TagTSCache.add add_cache key z;z
449
450 let child_sibling_labels tree =
451   let table_c = Array.create (tree_num_tags tree) TagS.empty in
452   let table_n = Array.copy table_c in
453   let rec loop node =
454     if node == nil then TagS.empty
455     else
456       let children = loop (tree_first_child tree node) in
457       let tag = tree_tag tree node in
458       let () =
459         let tc = table_c.(tag) in
460         if _subset children tc then ()
461         else table_c.(tag) <-  _union tc children
462       in
463       let siblings = loop (tree_next_sibling tree node) in
464       let () =
465         let tn = table_n.(tag) in
466         if _subset siblings tn then ()
467         else table_n.(tag) <- _union tn siblings
468       in
469       _add tag siblings
470   in
471   ignore (loop root);
472   table_c, table_n
473
474 let descendant_labels tree =
475   let table_d = Array.create (tree_num_tags tree) TagS.empty in
476   let rec loop node =
477     if node == nil then  TagS.empty else
478       let d1 = loop (tree_first_child tree node) in
479       let d2 = loop (tree_next_sibling tree node) in
480       let tag = tree_tag tree node in
481       let () =
482         let td = table_d.(tag) in
483         if _subset d1 td then ()
484         else table_d.(tag) <- _union td d1;
485       in
486       _add tag (_union d1 d2)
487   in
488   ignore (loop root);
489   table_d
490
491 let collect_labels tree =
492   let table_f = Array.create (tree_num_tags tree) TagS.empty in
493   let table_n = Array.copy table_f in
494   let table_c = Array.copy table_f in
495   let table_d = Array.copy table_f in
496   let rec loop node foll_siblings descendants followings =
497     if node == nil then foll_siblings, descendants, followings else
498       let tag = tree_tag tree node in
499       let () =
500         let tf = table_f.(tag) in
501         if _subset followings tf then ()
502         else table_f.(tag) <- _union tf followings in
503       let () =
504         let tn = table_n.(tag) in
505         if _subset foll_siblings tn then ()
506         else table_n.(tag) <- _union tn foll_siblings in
507       let children, n_descendants, n_followings =
508         loop (tree_last_child tree node) TagS.empty TagS.empty followings
509       in
510       let () =
511         let tc = table_c.(tag) in
512         if _subset children tc then ()
513         else table_c.(tag) <- _union tc children
514       in
515       let () =
516         let td = table_d.(tag) in
517         if _subset n_descendants td then ()
518         else table_d.(tag) <- _union td n_descendants
519       in
520       loop (tree_prev_sibling tree node)
521         (_add tag foll_siblings)
522         (_add tag (_union n_descendants descendants))
523         (_add tag n_followings)
524   in
525   ignore (loop root TagS.empty TagS.empty TagS.empty);
526   table_f, table_n, table_c, table_d
527
528
529 let is_nil t = t == nil
530 let is_node t = t != nil
531 let is_root t = t == root
532
533 let node_of_t t  =
534   Logger.print err_formatter "Initializing tag structure@\n";
535   let _ = Tag.init (mk_tag_ops t) in
536   Logger.print err_formatter "Starting tag table construction@\n";
537   let f, n, c, d = time collect_labels t ~msg:"Building tag relationship table" in
538   let c = Array.map TagS.to_ptset c in
539   let n = Array.map TagS.to_ptset n in
540   let f = Array.map TagS.to_ptset f in
541   let d = Array.map TagS.to_ptset d in
542   let () = clear_cache () in
543   let attributes = Ptset.Int.add Tag.attribute d.(Tag.attribute) in
544   let elements = Ptset.Int.add Tag.document_node
545     (Ptset.Int.remove Tag.pcdata
546        (Ptset.Int.diff d.(Tag.document_node) attributes))
547   in
548   { doc= t;
549     attributes = attributes;
550     attribute_array = Array.of_list (Ptset.Int.elements attributes);
551     elements = elements;
552     children = c;
553     siblings = n;
554     descendants = d;
555     followings = f
556
557   }
558
559
560 let parse_xml_uri str = node_of_t (TreeBuilder.parse_file str)
561 let parse_xml_string str = node_of_t (TreeBuilder.parse_string str)
562
563 let size t = tree_size t.doc;;
564
565 let magic_string = "SXSI_INDEX"
566 let version_string = "3"
567
568 let pos fd =
569   Unix.lseek fd 0  Unix.SEEK_CUR
570
571 let pr_pos fd = Logger.print err_formatter "At position %i@\n" (pos fd)
572
573 let write fd s =
574   let sl = String.length s in
575   let ssl = Printf.sprintf "%020i" sl in
576   ignore (Unix.write fd ssl 0 20);
577   ignore (Unix.write fd s 0 (String.length s))
578
579 let rec really_read fd buffer start length =
580   if length <= 0 then () else
581     match Unix.read fd buffer start length with
582       0 -> raise End_of_file
583     | r -> really_read fd buffer (start + r) (length - r);;
584
585 let read fd =
586   let buffer = String.create 20 in
587   let _ =  really_read fd buffer 0 20 in
588   let size = int_of_string buffer in
589   let buffer = String.create size in
590   let _ =  really_read fd buffer 0 size in
591   buffer
592
593 let save_tag_table channel t =
594   let t = Array.map (fun s -> Array.of_list (Ptset.Int.elements s)) t in
595   Marshal.to_channel channel t []
596
597 let save t str =
598   let fd = Unix.openfile str [ Unix.O_WRONLY;Unix.O_TRUNC;Unix.O_CREAT] 0o644 in
599   let out_c = Unix.out_channel_of_descr fd in
600   let index_prefix = Filename.chop_suffix str ".srx" in
601   let _ = set_binary_mode_out out_c true in
602   output_string out_c magic_string;
603   output_char out_c '\n';
604   output_string out_c version_string;
605   output_char out_c '\n';
606   save_tag_table out_c t.children;
607   save_tag_table out_c t.siblings;
608   save_tag_table out_c t.descendants;
609   save_tag_table out_c t.followings;
610     (* we need to move the fd to the correct position *)
611   flush out_c;
612   ignore (Unix.lseek fd (pos_out out_c) Unix.SEEK_SET);
613   tree_save t.doc fd index_prefix;
614   close_out out_c
615 ;;
616 let load_tag_table channel =
617   let table : int array array = Marshal.from_channel channel in
618   Array.map (fun a -> Ptset.Int.from_list (Array.to_list a)) table
619
620 let load ?(sample=64) ?(load_text=true) str =
621   let fd = Unix.openfile str [ Unix.O_RDONLY ] 0o644 in
622   let in_c = Unix.in_channel_of_descr fd in
623   let index_prefix = Filename.chop_suffix str ".srx" in
624   let _ = set_binary_mode_in in_c true in
625   let load_table () =
626     (let ms = input_line in_c in if ms <> magic_string then failwith "Invalid index file");
627     (let vs = input_line in_c in if vs <> version_string then failwith "Invalid version file");
628     let c = load_tag_table in_c in
629     let s = load_tag_table in_c in
630     let d = load_tag_table in_c in
631     let f = load_tag_table in_c in
632     c,s,d,f
633   in
634   let c, s, d, f = time ~msg:"Loading tag table"(load_table) () in
635   ignore(Unix.lseek fd (pos_in in_c) Unix.SEEK_SET);
636   let xml_tree = tree_load fd index_prefix load_text sample in
637   let () = Tag.init (Obj.magic xml_tree) in
638   let attributes = Ptset.Int.add Tag.attribute d.(Tag.attribute) in
639   let elements = Ptset.Int.add Tag.document_node
640     (Ptset.Int.remove Tag.pcdata
641        (Ptset.Int.diff d.(Tag.document_node) attributes))
642   in
643   let tree = { doc = xml_tree;
644                attributes = attributes;
645                attribute_array = Array.of_list (Ptset.Int.elements attributes);
646                elements = elements;
647                children = c;
648                siblings = s;
649                descendants = d;
650                followings = f
651              }
652   in close_in in_c;
653   tree
654
655
656
657
658 let equal a b = a == b
659
660 let nts = function
661 -1 -> "Nil"
662   | i -> Printf.sprintf "Node (%i)"  i
663
664 let dump_node t = nts (Node.to_int t)
665
666
667
668 type query_result = { bv : bit_vector;
669                       pos : node array;
670                     }
671
672 external tree_flush : tree -> Unix.file_descr -> unit = "caml_xml_tree_flush"
673 let flush t fd = tree_flush t.doc fd
674
675 external text_prefix : tree -> string -> bool -> query_result = "caml_text_collection_prefix_bv"
676 let text_prefix t s b = text_prefix t.doc s b
677
678 external text_suffix : tree -> string -> bool -> query_result = "caml_text_collection_suffix_bv"
679 let text_suffix t s b = text_suffix t.doc s b
680
681 external text_equals : tree -> string -> bool -> query_result = "caml_text_collection_equals_bv"
682 let text_equals t s b = text_equals t.doc s b
683
684 external text_contains : tree -> string -> bool -> query_result = "caml_text_collection_contains_bv"
685 let text_contains t s b = text_contains t.doc s b
686
687
688 module Predicate = Hcons.Make (
689   struct
690     type _t = t
691     type t = (_t -> node -> bool) ref
692     let hash t = Hashtbl.hash t
693     let equal t1 t2 = t1 == t2
694   end)
695
696 let string_of_query query =
697   match query with
698   | `Prefix -> "starts-with"
699   | `Suffix -> "ends-with"
700   | `Equals -> "equals"
701   | `Contains -> "contains"
702 ;;
703
704 let query_fun = function
705   | `Prefix -> text_prefix
706   | `Suffix -> text_suffix
707   | `Equals -> text_equals
708   | `Contains -> text_contains
709 ;;
710
711 let _pred_cache = Hashtbl.create 17
712 ;;
713 let mk_pred query s =
714   let f = query_fun query  in
715   let memo = ref (fun _ _ -> failwith "Undefined") in
716   memo := begin fun tree node ->
717     let results =
718       try Hashtbl.find _pred_cache (query,s) with
719         Not_found ->
720           time ~count:1 ~msg:(Printf.sprintf "Computing text query %s(%s)"
721                                 (string_of_query query) s)
722             (f tree) s true
723     in
724     let bv = results.bv in
725     memo := begin fun _ n ->
726       bit_vector_unsafe_get bv (Node.to_int n)
727     end;
728     bit_vector_unsafe_get bv (Node.to_int node)
729   end;
730   Predicate.make memo
731
732
733 let full_text_prefix t s = (text_prefix t s true).pos
734
735 let full_text_suffix t s = (text_suffix t s true).pos
736
737 let full_text_equals t s = (text_equals t s true).pos
738
739 let full_text_contains t s = (text_contains t s true).pos
740
741 let full_text_query q t s =
742   let res = (query_fun q) t s true in
743   Hashtbl.replace _pred_cache (q,s) res;
744   res.pos