Skip to content
This repository
Fetching contributors…

Octocat-spinner-32-eaf2f5

Cannot retrieve contributors at this time

file 313 lines (272 sloc) 10.485 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312
(* Copyright 2003 Yamagata Yoriyuki. distributed with LGPL *)
(* Modified by Edgar Friendly <thelema314@gmail.com> *)

module Core = struct

  type 'a t = (int * int * 'a) BatAvlTree.tree

  include BatAvlTree

  let singleton n v = singleton_tree (n, n, v)

  let make eq l (n1, n2, v) r =
    let n1, l =
      if is_empty l || n1 = min_int then n1, empty else
let (k1, k2, v0), l' = split_rightmost l in
if k2 + 1 = n1 && eq v v0 then k1, l' else n1, l in
    let n2, r =
      if is_empty r || n2 = max_int then n2, empty else
let (k1, k2, v0), r' = split_leftmost r in
if n2 + 1 = k1 && eq v v0 then k2, r' else n2, r in
    make_tree l (n1, n2, v) r

  let rec add ?(eq = (==)) n v m =
    if is_empty m then make_tree empty (n, n, v) empty else
      let (n1, n2, v0) as x = root m in
      let l = left_branch m in
      let r = right_branch m in
      if n1 <> min_int && n = n1 - 1 && eq v v0 then
make eq l (n, n2, v) r
      else if n < n1 then
make_tree (add n v l) x r
      else if n1 <= n && n <= n2 then
if eq v v0 then m else
let l =
if n1 = n then l else
make_tree l (n1, n - 1, v0) empty in
let r =
if n2 = n then r else
make_tree empty (n + 1, n2, v0) r in
make eq l (n, n, v) r
      else if n2 <> max_int && n = n2 + 1 && eq v v0 then
make eq l (n1, n, v) r
      else
make_tree l x (add n v r)

  let rec from n s =
    if is_empty s then empty else
      let (n1, n2, v) as x = root s in
      let s0 = left_branch s in
      let s1 = right_branch s in
      if n < n1 then make_tree (from n s0) x s1 else
if n > n2 then from n s1 else
make_tree empty (n, n2, v) s1

  let after n s = if n = max_int then empty else from (n + 1) s

  let rec until n s =
    if is_empty s then empty else
      let (n1, n2, v) as x = root s in
      let s0 = left_branch s in
      let s1 = right_branch s in
      if n > n2 then make_tree s0 x (until n s1) else
if n < n1 then until n s0 else
make_tree s0 (n1, n, v) empty

  let before n s = if n = min_int then empty else until (n - 1) s

  let add_range ?(eq=(==)) n1 n2 v s =
    if n1 > n2 then invalid_arg "IMap.add_range" else
      make eq (before n1 s) (n1, n2, v) (after n2 s)

  let rec find n m =
    if is_empty m then raise Not_found else
      let (n1, n2, v) = root m in
      if n < n1 then find n (left_branch m) else
if n1 <= n && n <= n2 then v else
find n (right_branch m)

  let rec remove n m =
    if is_empty m then empty else
      let (n1, n2, v) as x = root m in
      let l = left_branch m in
      let r = right_branch m in
      if n < n1 then
make_tree (remove n l) x r
      else if n1 = n then
if n2 = n then concat l r else
make_tree l (n + 1, n2, v) r
      else if n1 < n && n < n2 then
make_tree (make_tree l (n1, n - 1, v) empty) (n + 1, n2, v) r
      else if n = n2 then
make_tree l (n1, n - 1, v) r
      else
make_tree l x (remove n r)

  let remove_range n1 n2 m =
    if n1 > n2 then invalid_arg "IMap.remove_range" else
      concat (before n1 m) (after n2 m)

  let rec mem n m =
    if is_empty m then false else
      let (n1, n2, _) = root m in
      if n < n1 then mem n (left_branch m) else
if n1 <= n && n <= n2 then true else
mem n (right_branch m)

  let iter_range proc m =
    BatAvlTree.iter (fun (n1, n2, v) -> proc n1 n2 v) m

  let fold_range f m a =
    BatAvlTree.fold (fun (n1, n2, v) a -> f n1 n2 v a) m a

  let fold f m a =
    let rec loop n1 n2 v a =
      let a = f n1 v a in
      if n1 = n2 then a else
loop (n1 + 1) n2 v a in
    fold_range loop m a

  let iter proc m =
    fold (fun n v () -> proc n v) m ()

  let rec map ?(eq=(=)) f m =
    if is_empty m then empty else
      let n1, n2, v = root m in
      let l = map ~eq f (left_branch m) in
      let r = map ~eq f (right_branch m) in
      let v = f v in
      make eq l (n1, n2, v) r

  let mapi ?eq f m = fold (fun n v a -> add ?eq n (f n v) a) m empty

  let rec map_range ?(eq=(=)) f m =
    if is_empty m then empty else
      let n1, n2, v = root m in
      let l = map_range ~eq f (left_branch m) in
      let r = map_range ~eq f (right_branch m) in
      let v = f n1 n2 v in
      make eq l (n1, n2, v) r


  let rec set_to_map s v =
    if is_empty s then empty else
      let (n1, n2) = root s in
      let l = left_branch s in
      let r = right_branch s in
      make_tree (set_to_map l v) (n1, n2, v) (set_to_map r v)

  let domain m =
    if is_empty m then empty else
      let (k1, k2, _), m' = split_leftmost m in
      let f n1 n2 _ (k1, k2, s) =
if n1 = k2 + 1 then (k1, n2, s) else
(n1, n2, make_tree s (k1, k2) empty) in
      let k1, k2, s = fold_range f m' (k1, k2, empty) in
      make_tree s (k1, k2) empty

  let map_to_set p m =
    let rec loop m =
      if is_empty m then None else
let (k1, k2, v), m' = split_leftmost m in
if p v then Some (k1, k2, m') else
loop m' in
    match loop m with
      Some (k1, k2, m') ->
      let f n1 n2 v (k1, k2, s) =
if p v then
if n1 = k2 + 1 then (k1, n2, s) else
(n1, n2, make_tree s (k1, k2) empty)
else
(k1, k2, s) in
      let (k1, k2, s) = fold_range f m' (k1, k2, empty) in
      make_tree s (k1, k2) empty
    | None -> empty

  module Enum = BatEnum

  (* Fold across two maps *)
  let fold2_range f m1 m2 acc =
    let e1 = enum m1 and e2 = enum m2 in
    let rec aux acc = function
None,None -> acc
      | Some (lo,hi,rx), None ->
aux (f lo hi (Some rx) None acc) (Enum.get e1, None)
      | None, Some (lo,hi,rx) ->
aux (f lo hi None (Some rx) acc) (None, Enum.get e2)
      | Some (lo1,hi1,rx1), Some (lo2,hi2,rx2) when lo1 < lo2 ->
let hi, v1 =
if hi1 > lo2 then lo2-1, Some (lo2,hi1,rx1)
else if hi1 = lo2 then hi1, Some (lo2,lo2,rx1)
else hi1, Enum.get e1
and v2 = Some (lo2,hi2,rx2) in
aux (f lo1 hi (Some rx1) None acc) (v1, v2)
      | Some (lo1,hi1,rx1), Some (lo2,hi2,rx2) when lo2 < lo1 ->
let hi, v2 =
if hi2 > lo1 then lo1-1, Some (lo1,hi2,rx2)
else if hi2 = lo1 then hi2, Some (lo1,lo1,rx2)
else hi2, Enum.get e2
and v1 = Some (lo1,hi1,rx1) in
aux (f lo2 hi None (Some rx2) acc) (v1,v2)
      | Some (lo1,hi1,rx1), Some (_lo2,hi2,rx2) (* lo1 = lo2 *) ->
let hi, v1, v2 =
if hi1 = hi2 then hi1, Enum.get e1, Enum.get e2
else if hi1 < hi2 then hi1, Enum.get e1, Some (hi1+1,hi2,rx2)
else (* hi2 < hi1 *) hi2, Some (hi2+1,hi1,rx1), Enum.get e2
in
(* printf "#@%a\n" print_rng (lo1, hi); *)
aux (f lo1 hi (Some rx1) (Some rx2) acc) (v1, v2)
    in
    aux acc (Enum.get e1, Enum.get e2)

  let union ~eq f m1 m2 =
    let insert lo hi v1 v2 m = match v1, v2 with
      | Some v1, Some v2 -> add_range ~eq lo hi (f v1 v2) m
      | Some x, None | None, Some x -> add_range ~eq lo hi x m
      | None, None -> assert false
    in
    fold2_range insert m1 m2 empty

  let merge ~eq f m1 m2 =
    let insert lo hi v1 v2 m =
      match f lo hi v1 v2 with None -> m | Some v -> add_range ~eq lo hi v m in
    fold2_range insert m1 m2 empty


  let forall2_range f m1 m2 =
    let e1 = enum m1 and e2 = enum m2 in
    let rec aux = function
None,None -> true
      | Some (lo,hi,rx), None ->
(f lo hi (Some rx) None) && aux (Enum.get e1, None)
      | None, Some (lo,hi,rx) ->
(f lo hi None (Some rx)) && aux (None, Enum.get e2)
      | Some (lo1,hi1,rx1), Some (lo2,hi2,rx2) when lo1 < lo2 ->
let hi, v1 =
if hi1 > lo2 then lo2-1, Some (lo2,hi1,rx1)
else hi1, Enum.get e1
and v2 = Some (lo2,hi2,rx2) in
(f lo1 hi (Some rx1) None) && aux (v1, v2)
      | Some (lo1,hi1,rx1), Some (lo2,hi2,rx2) when lo2 < lo1 ->
let hi, v2 =
if hi2 > lo1 then lo1-1, Some (lo1,hi2,rx2)
else hi2, Enum.get e2
and v1 = Some (lo1,hi1,rx1) in
(f lo2 hi None (Some rx2)) && aux (v1,v2)
      | Some (lo1,hi1,rx1), Some (_,hi2,rx2) (* lo1 = lo2 *) ->
let hi, v1, v2 =
if hi1 = hi2 then hi1, Enum.get e1, Enum.get e2
else if hi1 < hi2 then hi1, Enum.get e1, Some (hi1+1,hi2,rx2)
else (* hi2 < hi1 *) hi2, Some (hi2+1,hi1,rx1), Enum.get e2
in
(f lo1 hi (Some rx1) (Some rx2)) && aux (v1, v2)
    in
    aux (Enum.get e1, Enum.get e2)
end

type 'a t = {m: 'a Core.t; eq: 'a -> 'a -> bool}
type key = int

let empty ~eq = {m = Core.empty; eq}
(*$T empty
is_empty (empty ~eq:(=))
*)

let singleton ~eq x y = {m = Core.singleton x y; eq}
(*$T singleton
not (is_empty (singleton ~eq:(=) 1 'x'))
find 1 (singleton ~eq:(=) 1 'x') = 'x'
try ignore(find 0 (singleton ~eq:(=) 1 'x')); false with Not_found -> true
*)

let is_empty {m} = Core.is_empty m
let add x y {m;eq} = {m=Core.add ~eq x y m; eq}

(*$= add as a & ~cmp:(List.eq (Tuple3.eq Int.equal Int.equal Int.equal)) ~printer:(List.print (Tuple3.print Int.print Int.print Int.print) |> IO.to_string)
[(0,2,0)] (empty ~eq:(=) |> a 0 0 |> a 2 0 |> a 1 0 |> enum |> List.of_enum)
*)
(*$= add as a & ~cmp:(List.eq (Tuple3.eq Int.equal Int.equal String.equal)) ~printer:(List.print (Tuple3.print Int.print Int.print String.print) |> IO.to_string)
[(0,2,"foo")] \
(empty ~eq:(=) |> a 0 "foo" |> a 2 "foo" |> a 1 "foo" |> enum |> List.of_enum)
*)


let add_range lo hi y {m;eq} = {m=Core.add_range ~eq lo hi y m; eq}
let find x {m} = Core.find x m
let remove x {m;eq} = {m=Core.remove x m; eq}
let remove_range lo hi {m;eq} = {m=Core.remove_range lo hi m; eq}
let from x {m;eq} = {m=Core.from x m; eq}
let after x {m;eq} = {m=Core.after x m; eq}
let until x {m;eq} = {m=Core.until x m; eq}
let before x {m;eq} = {m=Core.before x m; eq}
let mem x {m} = Core.mem x m
let iter f {m} = Core.iter f m
let iter_range f {m} = Core.iter_range f m
let map ?(eq=(=)) f {m} = {m=Core.map ~eq f m; eq}
let mapi ?(eq=(=)) f {m} = {m=Core.mapi ~eq f m; eq}
let map_range ?(eq=(=)) f {m} = {m = Core.map_range ~eq f m; eq}
let fold f {m} x0 = Core.fold f m x0
let fold_range f {m} x0 = Core.fold_range f m x0
let set_to_map ?(eq=(=)) s x = {m = Core.set_to_map s x; eq}
let domain {m} = Core.domain m
let map_to_set f {m} = Core.map_to_set f m
let enum {m} = Core.enum m
let fold2_range f {m=m1} {m=m2} x0 = Core.fold2_range f m1 m2 x0
let union f {m=m1;eq} {m=m2} = {m=Core.union ~eq f m1 m2; eq}
let merge ?(eq=(=)) f {m=m1} {m=m2} = {m=Core.merge ~eq f m1 m2; eq}
let forall2_range f {m=m1} {m=m2} = Core.forall2_range f m1 m2

module Infix = struct
  let (-->) {m} k = Core.find k m
  let (<--) m (k,v) = add k v m
end
Something went wrong with that request. Please try again.