src/utils/lib/intmap.ml

   1 (***********************************************************************)
   2 (*                                                                     *)
   3 (*                           Objective Caml                            *)
   4 (*                                                                     *)
   5 (*            Xavier Leroy, projet Cristal, INRIA Rocquencourt         *)
   6 (*                                                                     *)
   7 (*  Copyright 1996 Institut National de Recherche en Informatique et   *)
   8 (*  en Automatique.  All rights reserved.  This file is distributed    *)
   9 (*  under the terms of the GNU Library General Public License, with    *)
  10 (*  the special exception on linking described in file ../LICENSE.     *)
  11 (*                                                                     *)
  12 (***********************************************************************)
  13
  14 type key = int
  15
  16 type 'a t =
  17   Empty
  18 | Node of 'a t * key * 'a * 'a t * int
  19
  20 let empty = Empty
  21
  22 let height = function
  23     Empty -> 0
  24   | Node(_,_,_,_,h) -> h
  25
  26 let create l x d r =
  27   let hl = height l and hr = height r in
  28   Node(l, x, d, r, (if hl >= hr then hl + 1 else hr + 1))
  29
  30 let bal l x d r =
  31   let hl = match l with Empty -> 0 | Node(_,_,_,_,h) -> h in
  32   let hr = match r with Empty -> 0 | Node(_,_,_,_,h) -> h in
  33   if hl > hr + 2 then begin
  34       match l with
  35         Empty -> invalid_arg "Map.bal"
  36       | Node(ll, lv, ld, lr, _) ->
  37           if height ll >= height lr then
  38             create ll lv ld (create lr x d r)
  39           else begin
  40               match lr with
  41                 Empty -> invalid_arg "Map.bal"
  42               | Node(lrl, lrv, lrd, lrr, _)->
  43                   create (create ll lv ld lrl) lrv lrd (create lrr x d r)
  44             end
  45     end else if hr > hl + 2 then begin
  46       match r with
  47         Empty -> invalid_arg "Map.bal"
  48       | Node(rl, rv, rd, rr, _) ->
  49           if height rr >= height rl then
  50             create (create l x d rl) rv rd rr
  51           else begin
  52               match rl with
  53                 Empty -> invalid_arg "Map.bal"
  54               | Node(rll, rlv, rld, rlr, _) ->
  55                   create (create l x d rll) rlv rld (create rlr rv rd rr)
  56             end
  57     end else
  58     Node(l, x, d, r, (if hl >= hr then hl + 1 else hr + 1))
  59
  60 let rec add x data = function
  61     Empty ->
  62       Node(Empty, x, data, Empty, 1)
  63   | Node(l, v, d, r, h) ->
  64       if x = v then
  65         Node(l, x, data, r, h)
  66       else if x < v then
  67         bal (add x data l) v d r
  68       else
  69         bal l v d (add x data r)
  70
  71 let rec find x = function
  72     Empty ->
  73       raise Not_found
  74   | Node(l, v, d, r, _) ->
  75       if x = v then d
  76       else find x (if x < v then l else r)
  77
  78 let rec mem x = function
  79     Empty ->
  80       false
  81   | Node(l, v, d, r, _) ->
  82       x = v || mem x (if x < v then l else r)
  83
  84 let rec merge t1 t2 =
  85   match (t1, t2) with
  86     (Empty, t) -> t
  87   | (t, Empty) -> t
  88   | (Node(l1, v1, d1, r1, h1), Node(l2, v2, d2, r2, h2)) ->
  89       bal l1 v1 d1 (bal (merge r1 l2) v2 d2 r2)
  90
  91 let rec remove x = function
  92     Empty ->
  93       Empty
  94   | Node(l, v, d, r, h) ->
  95       if x = v then
  96         merge l r
  97       else if x < v then
  98         bal (remove x l) v d r
  99       else
 100         bal l v d (remove x r)
 101
 102 let rec iter f = function
 103     Empty -> ()
 104   | Node(l, v, d, r, _) ->
 105       iter f l; f v d; iter f r
 106
 107 let rec map f = function
 108     Empty               -> Empty
 109   | Node(l, v, d, r, h) -> Node(map f l, v, f d, map f r, h)
 110
 111 let rec mapi f = function
 112     Empty               -> Empty
 113   | Node(l, v, d, r, h) -> Node(mapi f l, v, f v d, mapi f r, h)
 114
 115 let rec fold f m accu =
 116   match m with
 117     Empty -> accu
 118   | Node(l, v, d, r, _) ->
 119       fold f l (f v d (fold f r accu))
 120
 121
 122 let rec length_aux len = function
 123     Empty -> len
 124   | Node(l, v, d, r, _) ->
 125       length_aux (length_aux (1+len) l) r
 126
 127 let length map =
 128   length_aux 0 map
 129
 130 let top = function
 131     Empty -> raise Not_found
 132   | Node(l, v, d, r, _) -> v, d
 133
 134 let rec infix_nth map res =
 135   match res, map with
 136     (count, Some _), _ -> res
 137   | (count, None), Empty -> res
 138   | (count, None), Node(l, v, d, r, _) ->
 139       infix_nth r (match infix_nth l res with
 140           (count, Some _) as res -> res
 141         | (count, None) ->
 142             if count = 0 then (count, Some d)
 143             else (count - 1, None))
 144
 145 let nth map n =
 146   match infix_nth map (n, None) with
 147     (_, None) -> raise Not_found
 148   | (_, Some node) -> node
 149
 150 let to_list map =
 151   let list = ref [] in
 152   iter (fun _ v -> list := v :: !list) map;
 153   !list
 154