?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
Initial snarf.
[shack.git] / libmojave / stdlib / lm_set.ml
blob075bc3894d9b8c5d8369b4dec7ba6af39f5779a1
1 (*
2 * Build a set using a red-black tree.
3 * Every node in the tree is colored either black or red.
4 * A red-black tree has the following invariants:
5 * 1. Every leaf is colored black
6 * 2. All children of every red node are black.
7 * 3. Every path from the root to a leaf has the
8 * same number of black nodes as every other path.
9 * 4. The root is always black.
10 *
11 * We get some corollaries:
12 * 1. The longest path from the root to a leaf is
13 * at most twice as long as the shortest path.
14 * 2. Both children of a red node are either leaves,
15 * or they are both not.
16 *
17 * This code is meant to be fast, so all the cases have
18 * been expanded, and the insert and delete functions are
19 * long (12 cases for insert, 18 for delete in lift_black).
20 *
21 * ----------------------------------------------------------------
22 *
23 * This file is part of MetaPRL, a modular, higher order
24 * logical framework that provides a logical programming
25 * environment for OCaml and other languages.
26 *
27 * See the file doc/htmlman/default.html or visit http://metaprl.org/
28 * for more information.
29 *
30 * Copyright (C) 1998-2005 PRL Group, Cornell University and Caltech
31 *
32 * This library is free software; you can redistribute it and/or
33 * modify it under the terms of the GNU Lesser General Public
34 * License as published by the Free Software Foundation,
35 * version 2.1 of the License.
36 *
37 * This library is distributed in the hope that it will be useful,
38 * but WITHOUT ANY WARRANTY; without even the implied warranty of
39 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
40 * Lesser General Public License for more details.
41 *
42 * You should have received a copy of the GNU Lesser General Public
43 * License along with this library; if not, write to the Free Software
44 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
45 *
46 * Additional permission is given to link this library with the
47 * OpenSSL project's "OpenSSL" library, and with the OCaml runtime,
48 * and you may distribute the linked executables. See the file
49 * LICENSE.libmojave for more details.
50 *
51 * Author: Jason Hickey
52 * jyh@cs.cornell.edu
53 *)
54 open Lm_printf
55 open Lm_set_sig
56
57 (*
58 * Make the set.
59 *)
60 module LmMake (Ord : OrderedType) =
61 struct
62 (************************************************************************
63 * TYPES *
64 ************************************************************************)
65
66 type elt = Ord.t
67
68 (*
69 * Table is a binary tree.
70 * Color is kept in the label to save space.
71 *)
72 (* %%MAGICBEGIN%% *)
73 type tree =
74 Leaf
75 | Red of elt * tree * tree * int
76 | Black of elt * tree * tree * int
77 (* %%MAGICEND%% *)
78
79 (*
80 * The tree is always balanced, so we don't need
81 * extra mutable fields.
82 *)
83 type t = tree
84
85 (*
86 * Path into the tree.
87 *)
88 type path =
89 Left of tree
90 | Right of tree
91 | Delete of tree
92
93 (*
94 * Exception for unchanged tree during insertion.
95 *)
96 exception Unchanged
97
98 (************************************************************************
99 * IMPLEMENTATION *
100 ************************************************************************)
101
102 (*
103 * Size of a table.
104 *)
105 let cardinality = function
106 Red (_, _, _, size)
107 | Black (_, _, _, size) ->
108 size
109 | Leaf ->
110 0
111
112 (*
113 * Add two nodes.
114 *)
115 let new_black key left right =
116 Black (key, left, right, cardinality left + cardinality right + 1)
117
118 let new_red key left right =
119 Red (key, left, right, cardinality left + cardinality right + 1)
120
121 (************************************************************************
122 * DEBUGGING *
123 ************************************************************************)
124
125 (*
126 * Print the tree.
127 *)
128 let rec pp_print_tree out tree =
129 match tree with
130 Black (_, left, right, size) ->
131 fprintf out "@[<v 3>Black(%d):@ %a@ %a@]" size pp_print_tree left pp_print_tree right
132 | Red (_, left, right, size) ->
133 fprintf out "@[<v 3>Red(%d):@ %a@ %a@]" size pp_print_tree left pp_print_tree right
134 | Leaf ->
135 fprintf out "Leaf"
136
137 let print_tree = pp_print_tree stdout
138
139 (*
140 * Check the size of the set.
141 *)
142 let check_size tree =
143 let abort tree' =
144 printf "%a@\n%a@\n" pp_print_tree tree pp_print_tree tree';
145 raise (Invalid_argument "check_size")
146 in
147 let rec check tree =
148 match tree with
149 Black (_, left, right, size) ->
150 if size <> check left + check right + 1 then
151 abort tree;
152 size
153
154 | Red (_, left, right, size) ->
155 if size <> check left + check right + 1 then
156 abort tree;
157 size
158
159 | Leaf ->
160 0
161 in
162 check tree
163
164 (*
165 * Check the red-invariant.
166 *)
167 let rec check_red = function
168 Red (_, Red _, _, _)
169 | Red (_, _, Red _, _) ->
170 raise (Failure "Lm_set.check_red")
171 | Red (_, left, right, _)
172 | Black (_, left, right, _) ->
173 check_red left;
174 check_red right
175 | Leaf ->
176 ()
177
178 (*
179 * Check the black invariant.
180 *)
181 let rec black_depth i = function
182 Black (_, left, _, _) ->
183 black_depth (succ i) left
184 | Red (_, left, _, _) ->
185 black_depth i left
186 | Leaf ->
187 i
188
189 let rec check_black_aux i j = function
190 Black (_, left, right, _) ->
191 check_black_aux i (succ j) left;
192 check_black_aux i (succ j) right
193 | Red (_, left, right, _) ->
194 check_black_aux i j left;
195 check_black_aux i j right
196 | Leaf ->
197 if j <> i then
198 raise (Failure "Lm_set.check_black")
199
200 let check_black tree =
201 check_black_aux (black_depth 0 tree) 0 tree
202
203 (*
204 * Check that all the nodes are sorted.
205 *)
206 let rec check_sort_lt key = function
207 Black (key', left, right, _)
208 | Red (key', left, right, _) ->
209 if Ord.compare key' key >= 0 then
210 raise (Failure "Lm_set.check_sort");
211 check_sort_lt key' left;
212 check_sort_gt_lt key' key right
213
214 | Leaf ->
215 ()
216
217 and check_sort_gt key = function
218 Black (key', left, right, _)
219 | Red (key', left, right, _) ->
220 if Ord.compare key' key <= 0 then
221 raise (Failure "Lm_set.check_sort");
222 check_sort_gt_lt key key' left;
223 check_sort_gt key right
224
225 | Leaf ->
226 ()
227
228 and check_sort_gt_lt key key' = function
229 Black (key'', left, right, _)
230 | Red (key'', left, right, _) ->
231 if Ord.compare key'' key <= 0 || Ord.compare key'' key' >= 0 then
232 raise (Failure "Lm_set.check_sort");
233 check_sort_gt_lt key key'' left;
234 check_sort_gt_lt key'' key' right
235
236 | Leaf ->
237 ()
238
239 let check_sort = function
240 Black (key, left, right, _) ->
241 check_sort_lt key left;
242 check_sort_gt key right
243 | Red _ ->
244 raise (Failure "Lm_set.check_sort: root is red")
245 | Leaf ->
246 ()
247
248 (*
249 * Perform all the checks.
250 *)
251 let check tree =
252 let _ =
253 check_red tree;
254 check_black tree;
255 check_sort tree;
256 check_size tree
257 in
258 tree
259
260 (************************************************************************
261 * INSERTION *
262 ************************************************************************)
263
264 (*
265 * Insert an entry into the tree.
266 *)
267 let rec insert key = function
268 Black (key0, left0, right0, size0) ->
269 begin
270 let comp = Ord.compare key key0 in
271 if comp = 0 then
272 raise Unchanged
273 else if comp < 0 then
274 match left0 with
275 Black _
276 | Leaf ->
277 (*
278 * Ok even if child becomes red.
279 *)
280 Black (key0, insert key left0, right0, succ size0)
281
282 | Red (key1, left1, right1, size1) ->
283 let comp = Ord.compare key key1 in
284 if comp = 0 then
285 raise Unchanged
286 else if comp < 0 then
287 match insert key left1, right0 with
288 Red _ as node, Red (key2, left2, right2, size2) ->
289 (*
290 * Recoloring:
291 *
292 * key0:b key0:r
293 * / \ / \
294 * key1:r key2:r key1:b key2:b
295 * / \ / \
296 * key2:r right1 key2:r right1
297 *)
298 Red (key0,
299 Black (key1, node, right1, succ size1),
300 Black (key2, left2, right2, size2),
301 succ size0)
302 | Red _ as node, _ ->
303 (*
304 * Rotation:
305 *
306 * key0:b key1:b
307 * / \ / \
308 * key1:r key2:b key3:r key0:b
309 * / \ / \
310 * key3:r right1 right1 key2:r
311 *)
312 Black (key1,
313 node,
314 new_red key0 right1 right0,
315 succ size0)
316 | node, _ ->
317 (*
318 * Inline:
319 *
320 * key0:b key0:b
321 * / \ / \
322 * key1:r key2 key1:r key2
323 * / \ / \
324 * key3:b right1 key3:b right1
325 *)
326 Black (key0,
327 new_red key1 node right1,
328 right0,
329 succ size0)
330 else
331 match insert key right1, right0 with
332 Red _ as node, Red (key2, left2, right2, size2) ->
333 (*
334 * Recoloring:
335 *
336 * key0:b key0:r
337 * / \ / \
338 * key1:r key2:r key1:b key2:b
339 * / \ / \
340 * left1 node:r left1 node:r
341 *)
342 Red (key0,
343 Black (key1, left1, node, succ size1),
344 Black (key2, left2, right2, size2),
345 succ size0)
346 | Red (key3, left3, right3, _), _ ->
347 (*
348 * Rotation:
349 *
350 * key0:b key3:b
351 * / \ / \
352 * key1:r right0 key1:r key0:r
353 * / \ / \ / \
354 * left1 key3:r left1 left3 right3 right0
355 * / \
356 * left3 right3
357 *)
358 Black (key3,
359 new_red key1 left1 left3,
360 new_red key0 right3 right0,
361 succ size0)
362 | node3, _ ->
363 (*
364 * Inline:
365 *
366 * key0:b
367 * / \
368 * key1:r right0
369 * / \
370 * left1 node3:b
371 *)
372 Black (key0,
373 new_red key1 left1 node3,
374 right0,
375 succ size0)
376 else
377 (* comp > 0 *)
378 match right0 with
379 Black _
380 | Leaf ->
381 (*
382 * Node can be replaced even if it becomes red.
383 *)
384 Black (key0, left0, insert key right0, succ size0)
385
386 | Red (key2, left2, right2, size2) ->
387 let comp = Ord.compare key key2 in
388 if comp = 0 then
389 raise Unchanged
390 else if comp < 0 then
391 match left0, insert key left2 with
392 Red (key1, left1, right1, size1), (Red _ as node) ->
393 (*
394 * Recoloring:
395 *
396 * key0:b key0:r
397 * / \ / \
398 * key1:r key2:r key1:b key2:b
399 * / \ / \
400 * node:r right2 node:r right2
401 *)
402 Red (key0,
403 Black (key1, left1, right1, size1),
404 Black (key2, node, right2, succ size2),
405 succ size0)
406 | _, Red (key3, left3, right3, _) ->
407 (*
408 * Rotate:
409 *
410 * key0:b key3:b
411 * / \ / \
412 * key1:b key2:r key0:r key2:r
413 * / \ / \ / \
414 * key3:r right2 left0 left3 right3 right2
415 * / \
416 * left3 right3
417 *)
418 Black (key3,
419 new_red key0 left0 left3,
420 new_red key2 right3 right2,
421 succ size0)
422 | _, node3 ->
423 (*
424 * Inline:
425 *
426 * key0:b
427 * / \
428 * left0 key2:r
429 * / \
430 * key3:b right2
431 *)
432 Black (key0,
433 left0,
434 new_red key2 node3 right2,
435 succ size0)
436 else
437 match left0, insert key right2 with
438 Red (key1, left1, right1, size1), (Red _ as node) ->
439 (*
440 * Recoloring:
441 *
442 * key0:b key0:r
443 * / \ / \
444 * key1:r key2:r key1:b key2:b
445 * / \ / \
446 * left2 node:r left2 node:r
447 *)
448 Red (key0,
449 Black (key1, left1, right1, size1),
450 Black (key2, left2, node, succ size2),
451 succ size0)
452 | _, (Red _ as node) ->
453 (*
454 * Rotation:
455 *
456 * key0:b key2:b
457 * / \ / \
458 * left0:b key2:r key0:r node:r
459 * / \ / \
460 * left2 node:r left0:b left2
461 *)
462 Black (key2,
463 new_red key0 left0 left2,
464 node,
465 succ size0)
466 | _, node3 ->
467 (*
468 * Inline:
469 *
470 * key0:b
471 * / \
472 * left0:b key2:r
473 * / \
474 * left2 node3:b
475 *)
476 Black (key0,
477 left0,
478 new_red key2 left2 node3,
479 succ size0)
480 end
481 | Leaf ->
482 (* Leaf is colored red *)
483 Red (key, Leaf, Leaf, 1)
484
485 | (Red _) ->
486 (* Red nodes will not come up *)
487 raise (Invalid_argument "Lm_set.insert")
488
489 (*
490 let insert key tree =
491 try insert key tree with
492 (Invalid_argument _) as exn ->
493 print tree;
494 print_newline ();
495 raise exn
496 *)
497
498 (*
499 * Add an element to the set.
500 *)
501 let add t key = match t with
502 Leaf ->
503 Black (key, Leaf, Leaf, 1)
504 | node ->
505 try
506 match insert key node with
507 Red (key, left, right, size) ->
508 Black (key, left, right, size)
509 | tree ->
510 tree
511 with
512 Unchanged ->
513 node
514
515 let add_list set keys =
516 List.fold_left add set keys
517
518 (************************************************************************
519 * REMOVAL *
520 ************************************************************************)
521
522 (*
523 * Construct a path during the removal.
524 *)
525 let rec delete key path node =
526 match node with
527 Black (key', left, right, _) ->
528 let comp = Ord.compare key key' in
529 if comp = 0 then
530 match left, right with
531 Leaf, Leaf ->
532 lift_black key path Leaf
533 | Red (key, left, right, size), Leaf ->
534 lift key path (Black (key, left, right, size))
535 | _ ->
536 delete_min (Delete node :: path) right
537 else if comp < 0 then
538 delete key (Left node :: path) left
539 else
540 delete key (Right node :: path) right
541 | Red (key', left, right, _) ->
542 let comp = Ord.compare key key' in
543 if comp = 0 then
544 match right with
545 Leaf ->
546 lift key path Leaf
547 | _ ->
548 delete_min (Delete node :: path) right
549 else if comp < 0 then
550 delete key (Left node :: path) left
551 else
552 delete key (Right node :: path) right
553 | Leaf ->
554 raise Not_found
555
556 and delete_min path node =
557 match node with
558 Black (key, Leaf, Leaf, _) ->
559 lift_black key path Leaf
560 | Black (key, Leaf, Red (key', left, right, size), _) ->
561 lift key path (Black (key', left, right, size))
562 | Red (key, Leaf, Leaf, _) ->
563 lift key path Leaf
564 | Black (_, left, _, _) ->
565 delete_min (Left node :: path) left
566 | Red (_, left, _, _) ->
567 delete_min (Left node :: path) left
568 | Leaf ->
569 raise Not_found
570
571 (*
572 * Copy the tree with no need to propagate black.
573 *)
574 and lift key path node =
575 match path, node with
576 Left (Black (key0, _, right0, size0)) :: path, left ->
577 lift key path (Black (key0, left, right0, pred size0))
578 | Left (Red (key0, _, right0, size0)) :: path, left ->
579 lift key path (Red (key0, left, right0, pred size0))
580 | Right (Black (key0, left0, _, size0)) :: path, right ->
581 lift key path (Black (key0, left0, right, pred size0))
582 | Right (Red (key0, left0, _, size0)) :: path, right ->
583 lift key path (Red (key0, left0, right, pred size0))
584 | Delete (Black (_, left0, _, size0)) :: path, right ->
585 lift key path (Black (key, left0, right, pred size0))
586 | Delete (Red (_, left0, _, size0)) :: path, right ->
587 lift key path (Red (key, left0, right, pred size0))
588 | [], node ->
589 node
590 | Left Leaf :: _, _
591 | Right Leaf :: _, _
592 | Delete Leaf :: _, _ ->
593 raise (Invalid_argument "lift")
594
595 (*
596 * Propagate the extra black up the tree.
597 *)
598 and lift_black key path node =
599 match path, node with
600 Left (Black (key0, _, right0, size0)) :: path, left ->
601 begin
602 match right0 with
603 Black (key2, left2, right2, size2) ->
604 begin
605 match left2, right2 with
606 _, Red (key3, left3, right3, size3) ->
607 (*
608 * key0:b key2:b
609 * / \ / \
610 * left:bb key2:b key0:b right2:b
611 * / \ / \
612 * left2 right2:r left:b left2
613 *)
614 lift key path (**)
615 (Black (key2,
616 new_black key0 left left2,
617 Black (key3, left3, right3, size3),
618 pred size0))
619
620 | Red (key3, left3, right3, _), _ ->
621 (*
622 * key0:b key3:b
623 * / \ / \
624 * left:bb key2:b key0:b key2:b
625 * / \ / \ / \
626 * key3:r right2:b left:b left3 right3 right2:b
627 * / \
628 * left3 right3
629 *)
630 lift key path (**)
631 (Black (key3,
632 new_black key0 left left3,
633 new_black key2 right3 right2,
634 pred size0))
635
636 | _ ->
637 (*
638 * key0:b key0:bb
639 * / \ / \
640 * left:bb key2:b left:b key2:r
641 * / \ / \
642 * left2:b right2:b left2:b right2:b
643 *)
644 lift_black key path (**)
645 (Black (key0,
646 left,
647 Red (key2, left2, right2, size2),
648 pred size0))
649 end
650
651 | Red (key2, left2, right2, _) ->
652 begin
653 match left2 with
654 Black (key3, Red (key4, left4, right4, _), d, _) ->
655 (*
656 * key0:b key2:b
657 * / \ / \
658 * left:bb key2:r key4:r right2:b
659 * / \ / \
660 * key3:b right2:b key0:b key3:b
661 * / \ / \ / \
662 * key4:r d left:b left4 right4 d
663 * / \
664 * left4 right4
665 *)
666 lift key path (**)
667 (Black (key2,
668 new_red key4 (**)
669 (new_black key0 left left4)
670 (new_black key3 right4 d),
671 right2,
672 pred size0))
673
674 | Black (key3, c, Red (key4, left4, right4, size4), _) ->
675 (*
676 * key0:b key2:b
677 * / \ / \
678 * left:bb key2:r key3:r right2
679 * / \ / \
680 * key3:b right2 key0:b key4:b
681 * / \ / \
682 * c key4:r left:b c
683 *)
684 lift key path (**)
685 (Black (key2,
686 new_red key3 (**)
687 (new_black key0 left c)
688 (Black (key4, left4, right4, size4)),
689 right2,
690 pred size0))
691
692 | Black (key3, c, d, _) ->
693 (*
694 * key0:b key2:b
695 * / \ / \
696 * left:bb key2:r key0:b right2:b
697 * / \ / \
698 * key3:b right2:b left:b key3:r
699 * / \ / \
700 * c:b d:b c:b d:b
701 *)
702 lift key path (**)
703 (Black (key2,
704 new_black key0 left (new_red key3 c d),
705 right2,
706 pred size0))
707
708 | Red _
709 | Leaf ->
710 raise (Invalid_argument "lift_black1")
711 end
712
713 | Leaf ->
714 raise (Invalid_argument "lift_black2")
715 end
716
717 | Right (Black (key0, left0, _, size0)) :: path, right ->
718 begin
719 match left0 with
720 Black (key1, left1, right1, size1) ->
721 begin
722 match left1, right1 with
723 Red (key3, left3, right3, size3), _ ->
724 (*
725 * key0:b key1:b
726 * / \ / \
727 * key1:b right:bb left1:b key0:b
728 * / \ / \
729 * left1:r right1 right1 right:b
730 *)
731 lift key path (**)
732 (Black (key1,
733 Black (key3, left3, right3, size3),
734 new_black key0 right1 right,
735 pred size0))
736
737 | _, Red (key3, left3, right3, _) ->
738 (*
739 * key0:b key3:b
740 * / \ / \
741 * key1:b right:bb key1:b key0:b
742 * / \ / \ / \
743 * left1:b key3:r left1:b left3 right3 right
744 * / \
745 * left3 right3
746 *)
747 lift key path (**)
748 (Black (key3,
749 new_black key1 left1 left3,
750 new_black key0 right3 right,
751 pred size0))
752
753 | _ ->
754 (*
755 * key0:b key0:bb
756 * / \ / \
757 * key1:b right:bb key1:r right:bb
758 * / \ / \
759 * left1:b right1:b left1:b right1:b
760 *)
761 lift_black key path (**)
762 (Black (key0,
763 Red (key1, left1, right1, size1),
764 right,
765 pred size0))
766
767 end
768
769 | Red (key1, left1, right1, _) ->
770 begin
771 match right1 with
772 Black (key3, d, Red (key4, left4, right4, _), _) ->
773 (*
774 * key0:b key1:b
775 * / \ / \
776 * key1:r right:bb left1:b key4:r
777 * / \ / \
778 * left1:b key3:b key3:b key0:b
779 * / \ / \ / \
780 * d key4:r d left4 right4 right:b
781 * / \
782 * left4 right4
783 *)
784 lift key path (**)
785 (Black (key1,
786 left1,
787 new_red key4 (**)
788 (new_black key3 d left4)
789 (new_black key0 right4 right),
790 pred size0))
791
792 | Black (key3, Red (key4, left4, right4, size4), c, _) ->
793 (*
794 * key0:b key1:b
795 * / \ / \
796 * key1:r right:bb left1 key3:r
797 * / \ / \
798 * left1 key3:b key4:b key0:b
799 * / \ / \
800 * key4:r c c right:b
801 *)
802 lift key path (**)
803 (Black (key1,
804 left1,
805 new_red key3 (**)
806 (Black (key4, left4, right4, size4))
807 (new_black key0 c right),
808 pred size0))
809
810 | Black (key3, c, d, size3) ->
811 (*
812 * key0:b key1:b
813 * / \ / \
814 * key1:r right:bb left1 key0:b
815 * / \ / \
816 * left1 key3:b key3:r right:b
817 * / \ / \
818 * c:b d:b c:b d:b
819 *)
820 lift key path (**)
821 (Black (key1,
822 left1,
823 new_black key0 (Red (key3, c, d, size3)) right,
824 pred size0))
825
826 | Red _
827 | Leaf ->
828 raise (Invalid_argument "lift_black3")
829 end
830
831 | Leaf ->
832 raise (Invalid_argument "lift_black4")
833 end
834
835 | Left (Red (key0, _, right0, size0)) :: path, left ->
836 begin
837 match right0 with
838 Black (key2, left2, right2, size2) ->
839 begin
840 match left2, right2 with
841 _, Red (key3, left3, right3, size3) ->
842 (*
843 * key0:r key2:r
844 * / \ / \
845 * left:bb key2:b key0:b right2:b
846 * / \ / \
847 * left2:b right2:r left:b left2:b
848 *)
849 lift key path (**)
850 (Red (key2,
851 new_black key0 left left2,
852 Black (key3, left3, right3, size3),
853 pred size0))
854
855 | Red (key3, left3, right3, _), _ ->
856 (*
857 * key0:r key3:b
858 * / \ / \
859 * left:bb key2:b key0:r key2:r
860 * / \ / \ / \
861 * key3:r right2 left:b left3 right3 right2
862 * / \
863 * left3 right3
864 *)
865 lift key path (**)
866 (Black (key3,
867 new_red key0 left left3,
868 new_red key2 right3 right2,
869 pred size0))
870
871 | _ ->
872 (*
873 * key0:r key0:b
874 * / \ / \
875 * left:bb key2:b left:b key2:r
876 * / \ / \
877 * left2:b right2:b left2:b right2:b
878 *)
879 lift key path (**)
880 (Black (key0,
881 left,
882 Red (key2, left2, right2, size2),
883 pred size0))
884 end
885 | Red _
886 | Leaf ->
887 raise (Invalid_argument "lift_black5")
888 end
889
890 | Right (Red (key0, left0, _, size0)) :: path, right ->
891 begin
892 match left0 with
893 Black (key1, left1, right1, size1) ->
894 begin
895 match left1, right1 with
896 Red (key3, left3, right3, size3), _ ->
897 (*
898 * key0:r key1:r
899 * / \ / \
900 * key1:b right:bb left1:b key0:b
901 * / \ / \
902 * left1:r right1 right1 right:b
903 *)
904 lift key path (**)
905 (Red (key1,
906 Black (key3, left3, right3, size3),
907 new_black key0 right1 right,
908 pred size0))
909
910 | _, Red (key3, left3, right3, _) ->
911 (*
912 * key0:r key3:b
913 * / \ / \
914 * key1:b right:bb key1:r key0:r
915 * / \ / \ / \
916 * left1 key3:r left1 left3 right3 right:b
917 * / \
918 * left3 right3
919 *)
920 lift key path (**)
921 (Black (key3,
922 new_red key1 left1 left3,
923 new_red key0 right3 right,
924 pred size0))
925
926 | _ ->
927 (*
928 * key0:r key0:b
929 * / \ / \
930 * key1:b right:bb key1:r right:b
931 * / \ / \
932 * left1:b right1:b left1:b right1:b
933 *)
934 lift key path (**)
935 (Black (key0,
936 Red (key1, left1, right1, size1),
937 right,
938 pred size0))
939
940 end
941 | Red _
942 | Leaf ->
943 raise (Invalid_argument "lift_black6")
944 end
945
946 | Delete (Black (_, left0, right0, size0)) :: path, node ->
947 lift_black key (Right (Black (key, left0, right0, size0)) :: path) node
948
949 | Delete (Red (_, left0, right0, size0)) :: path, node ->
950 lift_black key (Right (Red (key, left0, right0, size0)) :: path) node
951
952 | [], node ->
953 node
954
955 | Left Leaf :: _, _
956 | Right Leaf :: _, _
957 | Delete Leaf :: _, _ ->
958 raise (Invalid_argument "lift_black7")
959
960 (*
961 * Remove the item.
962 *)
963 let remove tree key =
964 try delete key [] tree with
965 Not_found ->
966 tree
967
968 let subtract_list tree keys =
969 List.fold_left remove tree keys
970
971 (************************************************************************
972 * UNION & INTERSECTION *
973 ************************************************************************)
974
975 (*
976 * Get the elements of the list.
977 *)
978 let rec to_list_aux elements = function
979 Black (key, left, right, _)
980 | Red (key, left, right, _) ->
981 to_list_aux (key :: to_list_aux elements right) left
982 | Leaf ->
983 elements
984
985 let to_list = to_list_aux []
986
987 let elements = to_list
988
989 let rec reverse elements = function
990 h :: t ->
991 reverse (h :: elements) t
992 | [] ->
993 elements
994
995 let rec merge elements elements1 elements2 =
996 match elements1, elements2 with
997 key1 :: tl1, key2 :: tl2 ->
998 let comp = Ord.compare key1 key2 in
999 if comp = 0 then
1000 merge (key1 :: elements) tl1 tl2
1001 else if comp < 0 then
1002 merge (key1 :: elements) tl1 elements2
1003 else
1004 merge (key2 :: elements) elements1 tl2
1005 | _, [] ->
1006 reverse elements1 elements
1007 | [], _ ->
1008 reverse elements2 elements
1009
1010 (*
1011 * Log of a number.
1012 *)
1013 let rec log2 i x =
1014 if 1 lsl i >= x then
1015 i
1016 else
1017 log2 (succ i) x
1018
1019 (*
1020 * Build a set from a list.
1021 *)
1022 let rec log2 i j =
1023 if 1 lsl i >= j then
1024 i
1025 else
1026 log2 (succ i) j
1027
1028 let rec of_sorted_array depth max_depth elements off len =
1029 if len = 1 then
1030 if depth = max_depth then
1031 Red (elements.(off), Leaf, Leaf, 1)
1032 else
1033 Black (elements.(off), Leaf, Leaf, 1)
1034 else if len = 2 then
1035 Black (elements.(off + 1), Red (elements.(off), Leaf, Leaf, 1), Leaf, 2)
1036 else
1037 let len2 = len lsr 1 in
1038 Black (elements.(off + len2),
1039 of_sorted_array (succ depth) max_depth elements off len2,
1040 of_sorted_array (succ depth) max_depth elements (off + len2 + 1) (len - len2 - 1),
1041 len)
1042
1043 let of_sorted_list = function
1044 [] ->
1045 Leaf
1046 | [key] ->
1047 Black (key, Leaf, Leaf, 1)
1048 | elements ->
1049 let elements = Array.of_list elements in
1050 let length = Lm_array_util.distinct compare elements in
1051 let max_depth = pred (log2 1 (succ length)) in
1052 of_sorted_array 0 max_depth elements 0 length
1053
1054 (*
1055 * Convert to a list.
1056 *)
1057 let rec to_list_aux l = function
1058 Black (key, left, right, _)
1059 | Red (key, left, right, _) ->
1060 to_list_aux (key :: to_list_aux l right) left
1061 | Leaf ->
1062 l
1063
1064 let to_list t =
1065 to_list_aux [] t
1066
1067 (*
1068 * Union flattens the two trees,
1069 * merges them, then creates a new tree.
1070 *)
1071 let rec union_aux s1 = function
1072 Black (key, left, right, _)
1073 | Red (key, left, right, _) ->
1074 union_aux (add (union_aux s1 left) key) right
1075 | Leaf ->
1076 s1
1077
1078 let union s1 s2 =
1079 let size1 = cardinality s1 in
1080 let size2 = cardinality s2 in
1081 if size1 < size2 then
1082 union_aux s2 s1
1083 else
1084 union_aux s1 s2
1085
1086 (*
1087 * See if two sets intersect.
1088 *)
1089 let rec intersect_aux elems1 elems2 =
1090 match elems1, elems2 with
1091 elem1 :: elems1', elem2 :: elems2' ->
1092 let comp = Ord.compare elem1 elem2 in
1093 if comp = 0 then
1094 true
1095 else if comp < 0 then
1096 intersect_aux elems1' elems2
1097 else
1098 intersect_aux elems1 elems2'
1099 | [], _
1100 | _, [] ->
1101 false
1102
1103 let intersectp s1 s2 =
1104 intersect_aux (to_list s1) (to_list s2)
1105
1106 (************************************************************************
1107 * IMPLEMENTATION *
1108 ************************************************************************)
1109
1110 (*
1111 * Search without reorganizing the tree.
1112 *)
1113 let rec mem t key = match t with
1114 Black (key', left, right, _)
1115 | Red (key', left, right, _) ->
1116 let comp = Ord.compare key key' in
1117 if comp = 0 then
1118 true
1119 else if comp < 0 then
1120 mem left key
1121 else
1122 mem right key
1123
1124 | Leaf ->
1125 false
1126
1127 (*
1128 * An empty tree is just a leaf.
1129 *)
1130 let empty = Leaf
1131
1132 let is_empty = function
1133 Leaf ->
1134 true
1135 | _ ->
1136 false
1137
1138 let singleton key =
1139 Black (key, Leaf, Leaf, 1)
1140
1141 let of_list l =
1142 List.fold_left (fun set item -> add set item) empty l
1143
1144 (*
1145 * Iterate a function over the hashtable.
1146 *)
1147 let rec iter f = function
1148 Black (key, left, right, _)
1149 | Red (key, left, right, _) ->
1150 iter f left;
1151 f key;
1152 iter f right
1153 | Leaf ->
1154 ()
1155
1156 (*
1157 * Fold a function over the subrange of the set
1158 *)
1159 let rec range_fold range f arg = function
1160 Black (key, left, right, _)
1161 | Red (key, left, right, _) ->
1162 let c = range key in
1163 if c > 0 then
1164 range_fold range f arg right
1165 else if c < 0 then
1166 range_fold range f arg left
1167 else
1168 let arg = range_fold range f arg left in
1169 let arg = f arg key in
1170 range_fold range f arg right
1171 | Leaf ->
1172 arg
1173
1174 (*
1175 * Fold a function over the set.
1176 *)
1177 let rec fold f arg = function
1178 Black (key, left, right, _)
1179 | Red (key, left, right, _) ->
1180 let arg = fold f arg left in
1181 let arg = f arg key in
1182 fold f arg right
1183 | Leaf ->
1184 arg
1185
1186 (*
1187 * Equality of sets.
1188 *)
1189 let rec equal set1 set2 =
1190 if cardinality set1 = cardinality set2 then
1191 let list1 = to_list set1 in
1192 let list2 = to_list set2 in
1193 List.for_all2 (fun x y -> Ord.compare x y = 0) list1 list2
1194 else
1195 false
1196
1197 (*
1198 * BUG: these functions are too slow!
1199 * Could be much more optimized.
1200 *)
1201 let filter pred s =
1202 fold (fun s' x ->
1203 if pred x then
1204 add s' x
1205 else
1206 s') empty s
1207
1208 let inter s1 s2 =
1209 let size1 = cardinality s1 in
1210 let size2 = cardinality s2 in
1211 let s1, s2 =
1212 if size1 < size2 then
1213 s1, s2
1214 else
1215 s2, s1
1216 in
1217 fold (fun s3 x ->
1218 if mem s2 x then
1219 add s3 x
1220 else
1221 s3) empty s1
1222
1223 let partition pred s =
1224 fold (fun (s1, s2) x ->
1225 if pred x then
1226 add s1 x, s2
1227 else
1228 s1, add s2 x) (empty, empty) s
1229
1230 let rec diff s = function
1231 Black (key, left, right, _)
1232 | Red (key, left, right, _) ->
1233 let s = remove s key in
1234 let s = diff s left in
1235 diff s right
1236 | Leaf ->
1237 s
1238
1239 let rec subset s1 s2 =
1240 match s1 with
1241 Black (key, left, right, _)
1242 | Red (key, left, right, _) ->
1243 mem s2 key && subset left s2 && subset right s2
1244 | Leaf ->
1245 true
1246
1247 let is_subset = subset
1248
1249 let compare s1 s2 =
1250 let rec compare s1 s2 =
1251 match s1, s2 with
1252 x1 :: s1, x2 :: s2 ->
1253 let cmp = Ord.compare x1 x2 in
1254 if cmp = 0 then
1255 compare s1 s2
1256 else
1257 cmp
1258 | [], [] ->
1259 0
1260 | [], _ :: _ ->
1261 -1
1262 | _ :: _, [] ->
1263 1
1264 in
1265 compare (to_list s1) (to_list s2)
1266
1267 (*
1268 * Choice.
1269 *)
1270 let rec min_elt = function
1271 Black (key, Leaf, _, _)
1272 | Red (key, Leaf, _, _) ->
1273 key
1274 | Black (_, left, _, _)
1275 | Red (_, left, _, _) ->
1276 min_elt left
1277 | Leaf ->
1278 raise Not_found
1279
1280 let rec max_elt = function
1281 Black (key, _, Leaf, _)
1282 | Red (key, _, Leaf, _) ->
1283 key
1284 | Black (_, _, right, _)
1285 | Red (_, _, right, _) ->
1286 max_elt right
1287 | Leaf ->
1288 raise Not_found
1289
1290 let choose = min_elt
1291
1292 (*
1293 * Predicates.
1294 *)
1295 let rec for_all pred = function
1296 Black (key, left, right, _)
1297 | Red (key, left, right, _) ->
1298 pred key && for_all pred left && for_all pred right
1299 | Leaf ->
1300 true
1301
1302 let rec exists pred = function
1303 Black (key, left, right, _)
1304 | Red (key, left, right, _) ->
1305 pred key || exists pred left || exists pred right
1306 | Leaf ->
1307 false
1308
1309 (*
1310 * Width.
1311 *)
1312 let cardinal = cardinality
1313
1314 (*
1315 * Filtering operations.
1316 *)
1317 let rec mem_filt s = function
1318 [] ->
1319 []
1320 | (h :: t) as l ->
1321 if mem s h then
1322 let rem = mem_filt s t in
1323 if rem == t then
1324 l
1325 else
1326 h :: rem
1327 else
1328 mem_filt s t
1329
1330 let rec not_mem_filt s = function
1331 [] ->
1332 []
1333 | (h :: t) as l ->
1334 if mem s h then
1335 not_mem_filt s t
1336 else
1337 let rem = not_mem_filt s t in
1338 if rem == t then
1339 l
1340 else
1341 h :: rem
1342
1343 let rec fst_mem_filt s = function
1344 [] ->
1345 []
1346 | (((v, _) as h) :: t) as l ->
1347 if mem s v then
1348 let rem = fst_mem_filt s t in
1349 if rem == t then
1350 l
1351 else
1352 h :: rem
1353 else
1354 fst_mem_filt s t
1355 end
1356
1357 module LmMakeDebug (Ord : OrderedTypeDebug) =
1358 struct
1359 module XSet = LmMake (Ord)
1360
1361 include XSet
1362
1363 (*
1364 * Print the tree.
1365 *)
1366 let rec pp_print out tree =
1367 fprintf out "@ ";
1368 match tree with
1369 Black (key, left, right, size) ->
1370 fprintf out "(@[<hv 0>Black@ %a:%d %a %a)@]" (**)
1371 Ord.print key
1372 size
1373 pp_print left
1374 pp_print right
1375
1376 | Red (key, left, right, size) ->
1377 fprintf out "(@[<hv 0>Red@ %a:%d %a %a)@]" (**)
1378 Ord.print key
1379 size
1380 pp_print left
1381 pp_print right
1382
1383 | Leaf ->
1384 output_string out "Leaf"
1385
1386 let print = pp_print
1387 end
1388
1389 module Make (Ord : OrderedType) : S with type elt = Ord.t =
1390 struct
1391 module XSet = LmMake (Ord)
1392
1393 include XSet
1394
1395 let mem x s =
1396 XSet.mem s x
1397
1398 let add x s =
1399 XSet.add s x
1400
1401 let remove x s =
1402 XSet.remove s x
1403
1404 let fold f s x =
1405 XSet.fold (fun x y -> f y x) x s
1406
1407 let partition f s =
1408 fst (XSet.partition f s)
1409 end
1410
1411 (*
1412 * -*-
1413 * Local Variables:
1414 * Caml-master: "compile"
1415 * End:
1416 * -*-
1417 *)
Compiler