module Int = Core.Core_int
module BatAvlTree = struct
type 'a tree =
| Empty
| Node of 'a tree * 'a * 'a tree * int
let empty = Empty
let is_empty = function
| Empty -> true
| Node _ -> false
let singleton_tree x =
Node (Empty, x, Empty, 1)
let left_branch = function
| Empty -> raise Not_found
| Node (l, _, _, _) -> l
let right_branch = function
| Empty -> raise Not_found
| Node (_, _, r, _) -> r
let root = function
| Empty -> raise Not_found
| Node (_, v, _, _) -> v
let height = function
| Empty -> 0
| Node (_, _, _, h) -> h
let create l v r =
let h' = 1 + Int.max (height l) (height r) in
assert (abs (height l - height r ) < 2);
Node (l, v, r, h')
let bal l v r =
let hl = height l in
let hr = height r in
if hl >= hr + 2 then
match l with
| Empty -> assert false
| Node (ll, lv, lr, _) ->
if height ll >= height lr then
create ll lv (create lr v r)
else
match lr with
| Empty -> assert false
| Node (lrl, lrv, lrr, _) ->
create (create ll lv lrl) lrv (create lrr v r)
else if hr >= hl + 2 then
match r with
| Empty -> assert false
| Node (rl, rv, rr, _) ->
if height rr >= height rl then
create (create l v rl) rv rr
else
match rl with
| Empty -> assert false
| Node (rll, rlv, rlr, _) ->
create (create l v rll) rlv (create rlr rv rr)
else
create l v r
let rec add_left v = function
| Empty -> Node (Empty, v, Empty, 1)
| Node (l, v', r, _) -> bal (add_left v l) v' r
let rec add_right v = function
| Empty -> Node (Empty, v, Empty, 1)
| Node (l, v', r, _) -> bal l v' (add_right v r)
let rec make_tree l v r =
match l , r with
| Empty, _ -> add_left v r
| _, Empty -> add_right v l
| Node (ll, lv, lr, lh), Node (rl, rv, rr, rh) ->
if lh > rh + 1 then bal ll lv (make_tree lr v r) else
if rh > lh + 1 then bal (make_tree l v rl) rv rr else
create l v r
let rec split_leftmost = function
| Empty -> raise Not_found
| Node (Empty, v, r, _) -> (v, r)
| Node (l, v, r, _) ->
let v0, l' = split_leftmost l in
(v0, make_tree l' v r)
let rec split_rightmost = function
| Empty -> raise Not_found
| Node (l, v, Empty, _) -> (v, l)
| Node (l, v, r, _) ->
let v0, r' = split_rightmost r in
(v0, make_tree l v r')
let rec concat t1 t2 =
match t1, t2 with
| Empty, _ -> t2
| _, Empty -> t1
| Node (l1, v1, r1, h1), Node (l2, v2, r2, h2) ->
if h1 < h2 then
make_tree (concat t1 l2) v2 r2
else
make_tree l1 v1 (concat r1 t2)
let rec iter proc = function
| Empty -> ()
| Node (l, v, r, _) ->
iter proc l;
proc v;
iter proc r
let rec fold f t init =
match t with
| Empty -> init
| Node (l, v, r, _) ->
let x = fold f l init in
let x = f v x in
fold f r x
let rec to_stream =
let module S = CFStream_stream in
function
| Empty -> S.empty ()
| Node (l, v, r, _) ->
S.append
(S.append
(S.of_lazy (lazy (to_stream l)))
(S.singleton v))
(S.of_lazy (lazy (to_stream r)))
end
include BatAvlTree
type t = (int * int) tree
type elt = int
let rec mem s (n:int) =
if is_empty s then false else
let v1, v2 = root s in
if n < v1 then mem (left_branch s) n else
if v1 <= n && n <= v2 then true else
mem (right_branch s) n
let rec intersects_range s i j =
if i > j then raise (Invalid_argument "iset_intersects_range") ;
if is_empty s then false
else
let v1, v2 = root s in
if j < v1 then intersects_range (left_branch s) i j
else if v2 < i then intersects_range (right_branch s) i j
else true
let rec add s n =
if is_empty s then make_tree empty (n, n) empty else
let (v1, v2) as v = root s in
let s0 = left_branch s in
let s1 = right_branch s in
if v1 <> min_int && n < v1 - 1 then make_tree (add s0 n) v s1 else
if v2 <> max_int && n > v2 + 1 then make_tree s0 v (add s1 n) else
if n + 1 = v1 then
if not (is_empty s0) then
let (u1, u2), s0' = split_rightmost s0 in
if u2 <> max_int && u2 + 1 = n then
make_tree s0' (u1, v2) s1
else
make_tree s0 (n, v2) s1
else
make_tree s0 (n, v2) s1
else if v2 + 1 = n then
if not (is_empty s1) then
let (u1, u2), s1' = split_leftmost s1 in
if n <> max_int && n + 1 = u1 then
make_tree s0 (v1, u2) s1'
else
make_tree s0 (v1, n) s1
else
make_tree s0 (v1, n) s1
else s
let rec from s ~n =
if is_empty s then empty else
let (v1, v2) as v = root s in
let s0 = left_branch s in
let s1 = right_branch s in
if n < v1 then make_tree (from s0 ~n) v s1 else
if n > v2 then from s1 ~n else
make_tree empty (n, v2) s1
let after s ~n = if n = max_int then empty else from s ~n:(n + 1)
let rec until s ~n =
if is_empty s then empty else
let (v1, v2) as v = root s in
let s0 = left_branch s in
let s1 = right_branch s in
if n > v2 then make_tree s0 v (until s1 ~n) else
if n < v1 then until s0 ~n else
make_tree s0 (v1, n) empty
let before s ~n = if n = min_int then empty else until s ~n:(n - 1)
let add_range s n1 n2 =
if n1 > n2 then invalid_arg (Printf.sprintf "ISet.add_range - %d > %d" n1 n2) else
let n1, l =
if n1 = min_int then n1, empty else
let l = until s ~n:(n1 - 1) in
if is_empty l then n1, empty else
let (v1, v2), l' = split_rightmost l in
if v2 + 1 = n1 then v1, l' else n1, l in
let n2, r =
if n2 = max_int then n2, empty else
let r = from s ~n:(n2 + 1) in
if is_empty r then n2, empty else
let (v1, v2), r' = split_leftmost r in
if n2 + 1 = v1 then v2, r' else n2, r in
make_tree l (n1, n2) r
let singleton n = singleton_tree (n, n)
let rec remove s n =
if is_empty s then empty else
let (v1, v2) as v = root s in
let s1 = left_branch s in
let s2 = right_branch s in
if n < v1 then make_tree (remove s1 n) v s2
else if n = v1 then
if v1 = v2 then concat s1 s2 else
make_tree s1 (v1 + 1, v2) s2
else if n > v1 && n < v2 then
let s = make_tree s1 (v1, n - 1) empty in
make_tree s (n + 1, v2) s2
else if n = v2 then make_tree s1 (v1, v2 - 1) s2 else
make_tree s1 v (remove s2 n)
let remove_range s n1 n2 =
if n1 > n2 then invalid_arg "ISet.remove_range" else
concat (before s n1) (after s n2)
let rec union s1 s2 =
if is_empty s1 then s2 else
if is_empty s2 then s1 else
let s1, s2 = if height s1 > height s2 then s1, s2 else s2, s1 in
let n1, n2 = root s1 in
let l1 = left_branch s1 in
let r1 = right_branch s1 in
let l2 = before s2 ~n:n1 in
let r2 = after s2 n2 in
let n1, l =
if n1 = min_int then n1, empty else
let l = union l1 l2 in
if is_empty l then n1, l else
let (v1, v2), l' = split_rightmost l in
if v2 + 1 = n1 then v1, l' else n1, l in
let n2, r =
if n1 = max_int then n2, empty else
let r = union r1 r2 in
if is_empty r then n2, r else
let (v1, v2), r' = split_leftmost r in
if n2 + 1 = v1 then v2, r' else n2, r in
make_tree l (n1, n2) r
let rec inter s1 s2 =
if is_empty s1 then empty else
if is_empty s2 then empty else
let s1, s2 = if height s1 > height s2 then s1, s2 else s2, s1 in
let n1, n2 = root s1 in
let l1 = left_branch s1 in
let r1 = right_branch s1 in
let l2 = before s2 n1 in
let r2 = after s2 n2 in
let m = until (from s2 n1) n2 in
concat (concat (inter l1 l2) m) (inter r1 r2)
let rec compl_aux n1 n2 s =
if is_empty s then add_range empty n1 n2 else
let v1, v2 = root s in
let l = left_branch s in
let r = right_branch s in
let l = if v1 = min_int then empty else compl_aux n1 (v1 - 1) l in
let r = if v2 = max_int then empty else compl_aux (v2 + 1) n2 r in
concat l r
let compl s = compl_aux min_int max_int s
let diff s1 s2 = inter s1 (compl s2)
let rec compare_aux x1 x2 =
match x1, x2 with
[], [] -> 0
| `Set s :: rest, x ->
if is_empty s then compare_aux rest x2 else
let l = left_branch s in
let v = root s in
let r = right_branch s in
compare_aux (`Set l :: `Range v :: `Set r :: rest) x
| _x, `Set s :: rest ->
if is_empty s then compare_aux x1 rest else
let l = left_branch s in
let v = root s in
let r = right_branch s in
compare_aux x1 (`Set l :: `Range v :: `Set r :: rest)
| `Range ((v1, v2)) :: rest1, `Range ((v3, v4)) :: rest2 ->
let sgn = Int.compare v1 v3 in
if sgn <> 0 then sgn else
let sgn = Int.compare v2 v4 in
if sgn <> 0 then sgn else
compare_aux rest1 rest2
| [], _ -> ~-1
| _, [] -> 1
let compare s1 s2 = compare_aux [`Set s1] [`Set s2]
let equal s1 s2 = compare s1 s2 = 0
let rec subset s1 s2 =
if is_empty s1 then true else
if is_empty s2 then false else
let v1, v2 = root s2 in
let l2 = left_branch s2 in
let r2 = right_branch s2 in
let l1 = before s1 v1 in
let r1 = after s1 v2 in
(subset l1 l2) && (subset r1 r2)
let fold_range s ~init ~f = BatAvlTree.fold (fun (n1, n2) x -> f n1 n2 x) s init
let fold s ~init ~f =
let rec g n1 n2 a =
if n1 = n2 then f n1 a else
g (n1 + 1) n2 (f n1 a) in
fold_range ~f:g s ~init
let iter s ~f = fold s ~init:() ~f:(fun n () -> f n)
let iter_range s ~f = BatAvlTree.iter (fun (n1, n2) -> f n1 n2) s
let for_all s ~f =
let rec test_range n1 n2 =
if n1 = n2 then f n1 else
f n1 && test_range (n1 + 1) n2 in
let rec test_set s =
if is_empty s then true else
let n1, n2 = root s in
test_range n1 n2 &&
test_set (left_branch s) &&
test_set (right_branch s) in
test_set s
let exists s ~f =
let rec test_range n1 n2 =
if n1 = n2 then f n1 else
f n1 || test_range (n1 + 1) n2 in
let rec test_set s =
if is_empty s then false else
let n1, n2 = root s in
test_range n1 n2 ||
test_set (left_branch s) ||
test_set (right_branch s) in
test_set s
let filter_range p n1 n2 a =
let rec loop n1 n2 a = function
None ->
if n1 = n2 then
make_tree a (n1, n1) empty
else
loop (n1 + 1) n2 a (if p n1 then Some n1 else None)
| Some v1 as x ->
if n1 = n2 then make_tree a (v1, n1) empty else
if p n1 then
loop (n1 + 1) n2 a x
else
loop (n1 + 1) n2 (make_tree a (v1, n1 - 1) empty) None in
loop n1 n2 a None
let filter s ~f = fold_range s ~f:(filter_range f) ~init:empty
let partition_range p n1 n2 (a, b) =
let rec loop n1 n2 acc =
let acc =
let a, b, (v, n) = acc in
if p n1 = v then acc else
if v then
(make_tree a (n, n1) empty, b, (not v, n1))
else
(a, make_tree b (n, n1) empty, (not v, n1)) in
if n1 = n2 then
let a, b, (v, n) = acc in
if v then (make_tree a (n, n1) empty, b) else
(a, make_tree b (n, n1) empty)
else
loop (n1 + 1) n2 acc in
loop n1 n2 (a, b, (p n1, n1))
let partition s ~f = fold_range ~f:(partition_range f) s ~init:(empty, empty)
let cardinal s =
fold_range ~f:(fun n1 n2 c -> c + n2 - n1 + 1) s ~init:0
let rev_ranges s =
fold_range ~f:(fun n1 n2 a -> (n1, n2) :: a) s ~init:[]
let rec burst_range n1 n2 a =
if n1 = n2 then n1 :: a else
burst_range n1 (n2 - 1) (n2 :: a)
let elements s =
let f a (n1, n2) = burst_range n1 n2 a in
List.fold_left f [] (rev_ranges s)
let ranges s = List.rev (rev_ranges s)
let min_elt s =
let (n, _), _ = split_leftmost s in
n
let max_elt s =
let (_, n), _ = split_rightmost s in
n
let choose s = fst (root s)
let of_list l = List.fold_left (fun s (lo,hi) -> add_range s lo hi) empty l
let of_stream e = CFStream_stream.fold ~f:(fun s (lo,hi) -> add_range s lo hi) ~init:empty e