Implement reverse mapping from preorder to nodes.
[tatoo.git] / src / naive_tree.ml
1 (***********************************************************************)
2 (*                                                                     *)
3 (*                               TAToo                                 *)
4 (*                                                                     *)
5 (*                     Kim Nguyen, LRI UMR8623                         *)
6 (*                   Université Paris-Sud & CNRS                       *)
7 (*                                                                     *)
8 (*  Copyright 2010-2012 Université Paris-Sud and Centre National de la *)
9 (*  Recherche Scientifique. All rights reserved.  This file is         *)
10 (*  distributed under the terms of the GNU Lesser General Public       *)
11 (*  License, with the special exception on linking described in file   *)
12 (*  ../LICENSE.                                                        *)
13 (*                                                                     *)
14 (***********************************************************************)
15
16 type node = {
17   tag : QName.t;
18   preorder : int;
19   mutable kind : Tree.NodeKind.t;
20   mutable data : string;
21   mutable first_child : node;
22   mutable next_sibling : node;
23   mutable parent: node;
24 }
25
26
27
28 let rec nil = {
29   tag = QName.nil;
30   kind = Tree.NodeKind.Element;
31   preorder = -1;
32   data = "";
33   first_child = nil;
34   next_sibling = nil;
35   parent = nil;
36 }
37
38 let dummy_tag = QName.of_string "#dummy"
39 let rec dummy = {
40   tag = dummy_tag;
41   kind = Tree.NodeKind.Element;
42   preorder = -1;
43   data = "";
44   first_child = dummy;
45   next_sibling = dummy;
46   parent = dummy;
47 }
48
49
50 type t = {
51   root : node;
52   size : int;
53   by_preorder : node array;
54   (* TODO add other intersting stuff *)
55 }
56
57
58
59 module Parser =
60 struct
61
62   type context = {
63     mutable stack : node list;
64     text_buffer : Buffer.t;
65     mutable current_preorder : int;
66   }
67
68   let print_node_ptr fmt n =
69     Format.fprintf fmt "<%s>"
70       (if n == nil then "NIL" else
71         if n == dummy then "DUMMY" else
72           "NODE " ^  string_of_int n.preorder)
73
74   let debug_node fmt node =
75     Format.fprintf fmt "{ tag=%s; preorder=%i; data=%S; first_child=%a; next_sibling=%a; parent=%a }"
76       (QName.to_string node.tag)
77       node.preorder
78       node.data
79       print_node_ptr node.first_child
80       print_node_ptr node.next_sibling
81       print_node_ptr node.parent
82
83
84   let debug_ctx fmt ctx =
85     Format.fprintf fmt "Current context: { preorder = %i\n; stack = \n%a\n }\n-------------\n"
86       ctx.current_preorder
87       (Pretty.print_list ~sep:";\n" debug_node) ctx.stack
88
89
90   let push n ctx = ctx.stack <- n :: ctx.stack
91   let pop ctx =
92     match ctx.stack with
93       [] -> failwith "XML parse error"
94     | e :: l -> ctx.stack <- l; e
95
96   let top ctx = match ctx.stack with
97     | [] -> failwith "XML parse error"
98     | e :: _ -> e
99
100   let next ctx =
101     let i = ctx.current_preorder in
102     ctx.current_preorder <- 1 + i;
103     i
104
105   let is_left n = n.next_sibling == dummy
106
107
108   let text_string = QName.to_string QName.text
109   let comment_string = QName.to_string QName.comment
110
111
112   let rec start_element_handler parser_ ctx tag attr_list =
113     do_text parser_ ctx;
114     let parent = top ctx in
115     let n = { tag = QName.of_string tag;
116               kind = Tree.NodeKind.Element;
117               preorder = next ctx;
118               data = "";
119               first_child = dummy;
120               next_sibling = dummy;
121               parent = parent;
122             }
123     in
124     if parent.first_child == dummy then parent.first_child <- n
125     else parent.next_sibling <- n;
126     push n ctx;
127     List.iter (do_attribute parser_ ctx) attr_list
128
129   and do_attribute parser_ ctx (att, value) =
130     let att_tag = QName.to_string (QName.attribute (QName.of_string att)) in
131     start_element_handler parser_ ctx att_tag [];
132     let n = top ctx in
133     n.data <- value;
134     n.kind <- Tree.NodeKind.Attribute;
135     end_element_handler parser_ ctx att_tag
136
137   and consume_closing ctx n =
138     if n.next_sibling != dummy then
139       let _ = pop ctx in consume_closing ctx (top ctx)
140
141   and end_element_handler parser_ ctx _ =
142     do_text parser_ ctx;
143     let node = top ctx in
144     if node.first_child == dummy then node.first_child <- nil
145     else begin
146       node.next_sibling <- nil;
147       consume_closing ctx node
148     end
149
150   and do_text parser_ ctx =
151     if Buffer.length ctx.text_buffer != 0 then
152       let s = Buffer.contents ctx.text_buffer in
153       Buffer.clear  ctx.text_buffer;
154       start_element_handler parser_ ctx text_string [];
155       let node = top ctx in
156       node.data <- s;
157       node.kind <- Tree.NodeKind.Text;
158       end_element_handler parser_ ctx text_string
159
160   and comment_handler parser_ ctx s =
161     do_text parser_ ctx;
162     start_element_handler parser_ ctx comment_string [];
163     let node = top ctx in
164     node.data <- s;
165     node.kind <- Tree.NodeKind.Comment;
166     end_element_handler parser_ ctx comment_string
167
168   and processing_instruction_handler parser_ ctx tag data =
169     do_text parser_ ctx;
170     let pi = QName.to_string
171       (QName.processing_instruction (QName.of_string tag))
172     in
173     start_element_handler parser_ ctx pi [];
174     let node = top ctx in
175     node.data <- data;
176     node.kind <- Tree.NodeKind.ProcessingInstruction;
177     end_element_handler parser_ ctx pi
178
179
180   let character_data_handler _parser ctx text =
181     Buffer.add_string ctx.text_buffer text
182
183   let create_parser () =
184     let ctx = { text_buffer = Buffer.create 512;
185                 current_preorder = 0;
186                 stack = [] } in
187     let psr = Expat.parser_create ~encoding:None in
188     Expat.set_start_element_handler psr (start_element_handler psr ctx);
189     Expat.set_end_element_handler psr (end_element_handler psr ctx);
190     Expat.set_character_data_handler
191       psr (character_data_handler psr ctx);
192     Expat.set_comment_handler psr (comment_handler psr ctx);
193     Expat.set_processing_instruction_handler psr
194       (processing_instruction_handler psr ctx);
195     push { tag = QName.document;
196            preorder = next ctx;
197            kind = Tree.NodeKind.Document;
198            data = "";
199            first_child = dummy;
200            next_sibling = dummy;
201            parent = nil;
202          } ctx;
203     (psr,
204      fun () ->
205        let node = top ctx in
206        node.next_sibling <- nil;
207        consume_closing ctx node;
208        Expat.final psr;
209        let root = List.hd ctx.stack in
210        root.next_sibling <- nil;
211        let a = Array.create ctx.current_preorder nil in
212        let rec loop n =
213          if n != nil then
214            begin
215              a.(n.preorder) <- n;
216              loop n.first_child;
217              loop n.next_sibling;
218            end
219        in
220        loop root;
221        { root = root;
222          size = ctx.current_preorder;
223          by_preorder = a
224        }
225     )
226
227   let error e parser_ =
228     let msg = Printf.sprintf "%i.%i %s"
229       (Expat.get_current_line_number parser_)
230       (Expat.get_current_column_number parser_)
231       (Expat.xml_error_to_string e)
232     in
233     raise (Tree.Parse_error msg)
234
235   let parse_string s =
236     let parser_, finalize = create_parser () in
237     try
238       Expat.parse parser_ s;
239       finalize ()
240     with
241       Expat.Expat_error e -> error e parser_
242
243   let parse_file fd =
244     let buffer = String.create 4096 in
245     let parser_, finalize = create_parser () in
246     let rec loop () =
247       let read = input fd buffer 0 4096 in
248       if read != 0 then
249         let () = Expat.parse_sub parser_ buffer 0 read in
250         loop ()
251     in try
252          loop (); finalize ()
253       with
254         Expat.Expat_error e -> error e parser_
255
256 end
257
258
259 let load_xml_file = Parser.parse_file
260 let load_xml_string = Parser.parse_string
261
262 let output_escape_string out s =
263   for i = 0 to String.length s - 1 do
264     match s.[i] with
265     | '<' -> output_string out "&lt;"
266     | '>' -> output_string out "&gt;"
267     | '&' -> output_string out "&amp;"
268     | '"' -> output_string out "&quot;"
269     | '\'' -> output_string out "&apos;"
270     | c -> output_char out c
271   done
272
273
274 let rec print_attributes ?(sep=true) out tree_ node =
275   if (node.kind == Tree.NodeKind.Attribute) then
276     let tag = QName.to_string (QName.remove_prefix node.tag) in
277     if sep then output_char out ' ';
278     output_string out tag;
279     output_string out "=\"";
280     output_escape_string out node.data;
281     output_char out '\"';
282     print_attributes out tree_ node.next_sibling
283   else
284     node
285
286 let rec print_xml out tree_ node =
287   if node != nil then
288   let () =
289     let open Tree.NodeKind in
290     match node.kind with
291     | Node -> ()
292     | Text -> output_escape_string out node.data
293     | Element | Document ->
294         let tag = QName.to_string node.tag in
295         output_char out '<';
296         output_string out tag;
297         let fchild = print_attributes out tree_ node.first_child in
298         if fchild == nil then output_string out "/>"
299         else begin
300           output_char out '>';
301           print_xml out tree_ fchild;
302           output_string out "</";
303           output_string out tag;
304           output_char out '>'
305         end
306     | Attribute -> ignore (print_attributes ~sep:false out tree_ node)
307     | Comment ->
308         output_string out "<!--";
309         output_string out node.data;
310         output_string out "-->"
311     | ProcessingInstruction ->
312         output_string out "<?";
313         output_string out (QName.to_string (QName.remove_prefix node.tag));
314         output_char out ' ';
315         output_string out node.data;
316         output_string out "?>"
317   in
318   print_xml out tree_ node.next_sibling
319
320 let print_xml out tree_ node =
321   let nnode =  { node with next_sibling = nil } in print_xml out tree_ nnode
322
323 let root t = t.root
324 let size t = t.size
325 let first_child _ n = n.first_child
326 let next_sibling _ n = n.next_sibling
327 let parent _ n = n.parent
328 let tag _ n = n.tag
329 let data _ n = n.data
330 let kind _ n = n.kind
331 let preorder _ n = n.preorder
332 let by_preorder t i =
333  if i >= 0 && i < t.size then Array.unsafe_get t.by_preorder i
334  else let e = Invalid_argument "by_preorder" in raise e
335 let print_node fmt n = Parser.debug_node fmt n