Add fold_left/right functions to the set interface (iterate in
[tatoo.git] / src / ptset.ml
1 (* Original file: *)
2 (***********************************************************************)
3 (*                                                                     *)
4 (*  Copyright (C) Jean-Christophe Filliatre                            *)
5 (*                                                                     *)
6 (*  This software is free software; you can redistribute it and/or     *)
7 (*  modify it under the terms of the GNU Library General Public        *)
8 (*  License version 2.1, with the special exception on linking         *)
9 (*  described in file http://www.lri.fr/~filliatr/ftp/ocaml/ds/LICENSE *)
10 (*                                                                     *)
11 (*  This software is distributed in the hope that it will be useful,   *)
12 (*  but WITHOUT ANY WARRANTY; without even the implied warranty of     *)
13 (*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.               *)
14 (*                                                                     *)
15 (***********************************************************************)
16
17 (* Modified by Kim Nguyen *)
18 (* The Patricia trees are themselves deeply hash-consed. The module
19    provides a Make (and Weak) functor to build hash-consed patricia
20    trees whose elements are Abstract hash-consed values.
21 *)
22
23 INCLUDE "utils.ml"
24
25 include Ptset_sig
26
27 module type HConsBuilder =
28   functor (H : Common_sig.HashedType) -> Hcons.S with type data = H.t
29
30 module Builder (HCB : HConsBuilder) (H : Hcons.Abstract) :
31   S with type elt = H.t =
32 struct
33   type elt = H.t
34
35   type 'a set =
36     | Empty
37     | Leaf of elt
38     | Branch of int * int * 'a * 'a
39
40   module rec Node : Hcons.S with type data = Data.t = HCB(Data)
41                             and Data : Common_sig.HashedType with type t = Node.t set =
42   struct
43     type t =  Node.t set
44     let equal x y =
45       match x,y with
46       | Empty,Empty -> true
47       | Leaf k1, Leaf k2 ->  k1 == k2
48       | Branch(b1,i1,l1,r1), Branch(b2,i2,l2,r2) ->
49           b1 == b2 && i1 == i2 && (Node.equal l1 l2) && (Node.equal r1 r2)
50
51       | (Empty|Leaf _|Branch _), _  -> false
52
53     let hash = function
54     | Empty -> 0
55     | Leaf i -> HASHINT2 (PRIME1, Uid.to_int (H.uid i))
56     | Branch (b,i,l,r) ->
57         HASHINT4(b, i, Uid.to_int l.Node.id, Uid.to_int r.Node.id)
58   end
59
60   include Node
61
62   let empty = Node.make Empty
63
64   let is_empty s = s.Node.node == Empty
65
66   let branch p m l r = Node.make (Branch(p,m,l,r))
67
68   let leaf k = Node.make (Leaf k)
69
70                             (* To enforce the invariant that a branch contains two non empty
71                                sub-trees *)
72   let branch_ne p m t0 t1 =
73     if (is_empty t0) then t1
74     else if is_empty t1 then t0 else branch p m t0 t1
75
76                             (******** from here on, only use the smart constructors ************)
77
78   let singleton k = leaf k
79
80   let is_singleton n =
81     match n.Node.node with
82       | Leaf _ -> true
83       | Branch _ | Empty -> false
84
85   let mem (k:elt) n =
86     let kid = (H.uid k :> int) in
87     let rec loop n = match n.Node.node with
88     | Empty -> false
89     | Leaf j ->  k == j
90     | Branch (p, _, l, r) -> loop (if kid <= p then l else r)
91     in loop n
92
93   let rec min_elt n = match n.Node.node with
94   | Empty -> raise Not_found
95   | Leaf k -> k
96   | Branch (_,_,s,_) -> min_elt s
97
98   let rec max_elt n = match n.Node.node with
99   | Empty -> raise Not_found
100   | Leaf k -> k
101   | Branch (_,_,_,t) -> max_elt t
102
103   let elements s =
104     let rec elements_aux acc n = match n.Node.node with
105     | Empty -> acc
106     | Leaf k -> k :: acc
107     | Branch (_,_,l,r) -> elements_aux (elements_aux acc r) l
108     in
109     elements_aux [] s
110
111
112   let zero_bit k m = (k land m) == 0
113
114   let mask k m  = (k lor (m-1)) land (lnot m)
115
116   external int_of_bool : bool -> int = "%identity"
117
118   let hb32 v0 =
119     let v = v0 lor (v0 lsr 1) in
120     let v = v lor (v lsr 2) in
121     let v = v lor (v lsr 4) in
122     let v = v lor (v lsr 8) in
123     let v = v lor (v lsr 16) in
124     ((v + 1) lsr 1) + (int_of_bool (v0 == 0))
125
126   let hb64 v0 =
127     let v = v0 lor (v0 lsr 1) in
128     let v = v lor (v lsr 2) in
129     let v = v lor (v lsr 4) in
130     let v = v lor (v lsr 8) in
131     let v = v lor (v lsr 16) in
132     let v = v lor (v lsr 32) in
133     ((v + 1) lsr 1) + (int_of_bool (v0 == 0))
134
135
136   let branching_bit p0 p1 = hb64 (p0 lxor p1)
137
138   let join p0 t0 p1 t1 =
139     let m = branching_bit p0 p1  in
140     let msk = mask p0 m in
141     if zero_bit p0 m then
142       branch_ne msk m t0 t1
143     else
144       branch_ne msk m t1 t0
145
146   let match_prefix k p m = (mask k m) == p
147
148   let add k t =
149     let kid = Uid.to_int (H.uid k) in
150     let rec ins n = match n.Node.node with
151     | Empty -> leaf k
152     | Leaf j ->  if j == k then n else join kid (leaf k) (Uid.to_int (H.uid j)) n
153     | Branch (p,m,t0,t1)  ->
154         if match_prefix kid p m then
155           if zero_bit kid m then
156             branch_ne p m (ins t0) t1
157           else
158             branch_ne p m t0 (ins t1)
159         else
160           join kid (leaf k)  p n
161     in
162     ins t
163
164   let remove k t =
165     let kid = (H.uid k :> int) in
166     let rec rmv n = match n.Node.node with
167     | Empty -> empty
168     | Leaf j  -> if  k == j then empty else n
169     | Branch (p,m,t0,t1) ->
170         if match_prefix kid p m then
171           if zero_bit kid m then
172             branch_ne p m (rmv t0) t1
173           else
174             branch_ne p m t0 (rmv t1)
175         else
176           n
177     in
178     rmv t
179
180   (* should run in O(1) thanks to hash consing *)
181
182   let equal a b = a == b
183
184   let compare a b = (Uid.to_int (Node.uid a)) - (Uid.to_int (Node.uid b))
185
186   let rec merge s t =
187     if equal s t (* This is cheap thanks to hash-consing *)
188     then s
189     else
190       match s.Node.node, t.Node.node with
191       | Empty, _  -> t
192       | _, Empty  -> s
193       | Leaf k, _ -> add k t
194       | _, Leaf k -> add k s
195       | Branch (p,m,s0,s1), Branch (q,n,t0,t1) ->
196           if m == n && match_prefix q p m then
197             branch p  m  (merge s0 t0) (merge s1 t1)
198           else if m > n && match_prefix q p m then
199             if zero_bit q m then
200               branch_ne p m (merge s0 t) s1
201             else
202               branch_ne p m s0 (merge s1 t)
203           else if m < n && match_prefix p q n then
204             if zero_bit p n then
205               branch_ne q n (merge s t0) t1
206             else
207               branch_ne q n t0 (merge s t1)
208           else
209         (* The prefixes disagree. *)
210             join p s q t
211
212
213
214
215   let rec subset s1 s2 = (equal s1 s2) ||
216     match s1.Node.node, s2.Node.node with
217     | Empty, _ -> true
218     | _, Empty -> false
219     | Leaf k1, _ -> mem k1 s2
220     | Branch _, Leaf _ -> false
221     | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
222         if m1 == m2 && p1 == p2 then
223           subset l1 l2 && subset r1 r2
224         else if m1 < m2 && match_prefix p1 p2 m2 then
225           if zero_bit p1 m2 then
226             subset l1 l2 && subset r1 l2
227           else
228             subset l1 r2 && subset r1 r2
229         else
230           false
231
232
233   let union s1 s2 = merge s1 s2
234                             (* Todo replace with e Memo Module *)
235
236   let rec inter s1 s2 =
237     if equal s1 s2
238     then s1
239     else
240       match s1.Node.node, s2.Node.node with
241       | Empty, _ -> empty
242       | _, Empty -> empty
243       | Leaf k1, _ -> if mem k1 s2 then s1 else empty
244       | _, Leaf k2 -> if mem k2 s1 then s2 else empty
245       | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
246           if m1 == m2 && p1 == p2 then
247             merge (inter l1 l2)  (inter r1 r2)
248           else if m1 > m2 && match_prefix p2 p1 m1 then
249             inter (if zero_bit p2 m1 then l1 else r1) s2
250           else if m1 < m2 && match_prefix p1 p2 m2 then
251             inter s1 (if zero_bit p1 m2 then l2 else r2)
252           else
253             empty
254
255   let rec diff s1 s2 =
256     if equal s1 s2
257     then empty
258     else
259       match s1.Node.node, s2.Node.node with
260       | Empty, _ -> empty
261       | _, Empty -> s1
262       | Leaf k1, _ -> if mem k1 s2 then empty else s1
263       | _, Leaf k2 -> remove k2 s1
264       | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
265           if m1 == m2 && p1 == p2 then
266             merge (diff l1 l2) (diff r1 r2)
267           else if m1 > m2 && match_prefix p2 p1 m1 then
268             if zero_bit p2 m1 then
269               merge (diff l1 s2) r1
270             else
271               merge l1 (diff r1 s2)
272           else if m1 < m2 && match_prefix p1 p2 m2 then
273             if zero_bit p1 m2 then diff s1 l2 else diff s1 r2
274           else
275             s1
276
277
278   (*s All the following operations ([cardinal], [iter], [fold], [for_all],
279     [exists], [filter], [partition], [choose], [elements]) are
280     implemented as for any other kind of binary trees. *)
281
282   let rec cardinal n = match n.Node.node with
283   | Empty -> 0
284   | Leaf _ -> 1
285   | Branch (_,_,t0,t1) -> cardinal t0 + cardinal t1
286
287   let rec iter f n = match n.Node.node with
288   | Empty -> ()
289   | Leaf k -> f k
290   | Branch (_,_,t0,t1) -> iter f t0; iter f t1
291
292   let rec fold_left f s accu = match s.Node.node with
293   | Empty -> accu
294   | Leaf k -> f k accu
295   | Branch (_,_,t0,t1) -> fold_left f t1 (fold_left f t0 accu)
296
297   let rec fold_right f s accu = match s.Node.node with
298   | Empty -> accu
299   | Leaf k -> f k accu
300   | Branch (_,_,t0,t1) -> fold_right f t0 (fold_right f t1 accu)
301
302   let fold f s accu = fold_left f s accu
303
304   let rec for_all p n = match n.Node.node with
305   | Empty -> true
306   | Leaf k -> p k
307   | Branch (_,_,t0,t1) -> for_all p t0 && for_all p t1
308
309   let rec exists p n = match n.Node.node with
310   | Empty -> false
311   | Leaf k -> p k
312   | Branch (_,_,t0,t1) -> exists p t0 || exists p t1
313
314   let rec filter pr n = match n.Node.node with
315   | Empty -> empty
316   | Leaf k -> if pr k then n else empty
317   | Branch (p,m,t0,t1) -> let n0 = filter pr t0 in
318                           let n1 = filter pr t1 in
319                           branch_ne p m n0 n1
320
321   let partition p s =
322     let rec part (t,f as acc) n = match n.Node.node with
323     | Empty -> acc
324     | Leaf k -> if p k then (add k t, f) else (t, add k f)
325     | Branch (_,_,t0,t1) -> part (part acc t0) t1
326     in
327     part (empty, empty) s
328
329   let rec choose n = match n.Node.node with
330   | Empty -> raise Not_found
331   | Leaf k -> k
332   | Branch (_, _,t0,_) -> choose t0   (* we know that [t0] is non-empty *)
333
334
335   let split x s =
336     let coll k (l, b, r) =
337       if k < x then add k l, b, r
338       else if k > x then l, b, add k r
339       else l, true, r
340     in
341     fold coll s (empty, false, empty)
342
343   (*s Additional functions w.r.t to [Set.S]. *)
344
345   let rec intersect s1 s2 = (equal s1 s2) ||
346     match  s1.Node.node, s2.Node.node with
347     | Empty, _ -> false
348     | _, Empty -> false
349     | Leaf k1, _ -> mem k1 s2
350     | _, Leaf k2 -> mem k2 s1
351     | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
352         if m1 == m2 && p1 == p2 then
353           intersect l1 l2 || intersect r1 r2
354         else if m1 > m2 && match_prefix p2 p1 m1 then
355           intersect (if zero_bit p2 m1 then l1 else r1) s2
356         else if m1 < m2 && match_prefix p1 p2 m2 then
357           intersect s1 (if zero_bit p1 m2 then l2 else r2)
358         else
359           false
360
361
362   let from_list l = List.fold_left (fun acc e -> add e acc) empty l
363
364
365 end
366
367 module Make = Builder(Hcons.Make)
368 module Weak = Builder(Hcons.Weak)
369
370 module PosInt
371   =
372 struct
373   include Make(Hcons.PosInt)
374   let print ppf s =
375     Format.pp_print_string ppf "{ ";
376     iter (fun i -> Format.fprintf ppf "%i " i) s;
377     Format.pp_print_string ppf "}";
378     Format.pp_print_flush ppf ()
379 end