bundles/extlib/extlib-1.5.2/pMap.ml

   1 (*
   2  * PMap - Polymorphic maps
   3  * Copyright (C) 1996-2003 Xavier Leroy, Nicolas Cannasse, Markus Mottl
   4  *
   5  * This library is free software; you can redistribute it and/or
   6  * modify it under the terms of the GNU Lesser General Public
   7  * License as published by the Free Software Foundation; either
   8  * version 2.1 of the License, or (at your option) any later version,
   9  * with the special exception on linking described in file LICENSE.
  10  *
  11  * This library 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.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with this library; if not, write to the Free Software
  18  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  19  *)
  20
  21 type ('k, 'v) map =
  22   | Empty
  23   | Node of ('k, 'v) map * 'k * 'v * ('k, 'v) map * int
  24
  25 type ('k, 'v) t =
  26   {
  27     cmp : 'k -> 'k -> int;
  28     map : ('k, 'v) map;
  29   }
  30
  31 let height = function
  32   | Node (_, _, _, _, h) -> h
  33   | Empty -> 0
  34
  35 let make l k v r = Node (l, k, v, r, max (height l) (height r) + 1)
  36
  37 let bal l k v r =
  38   let hl = height l in
  39   let hr = height r in
  40   if hl > hr + 2 then
  41     match l with
  42     | Node (ll, lk, lv, lr, _) ->
  43         if height ll >= height lr then make ll lk lv (make lr k v r)
  44         else
  45           (match lr with
  46           | Node (lrl, lrk, lrv, lrr, _) ->
  47               make (make ll lk lv lrl) lrk lrv (make lrr k v r)
  48           | Empty -> assert false)
  49     | Empty -> assert false
  50   else if hr > hl + 2 then
  51     match r with
  52     | Node (rl, rk, rv, rr, _) ->
  53         if height rr >= height rl then make (make l k v rl) rk rv rr
  54         else
  55           (match rl with
  56           | Node (rll, rlk, rlv, rlr, _) ->
  57               make (make l k v rll) rlk rlv (make rlr rk rv rr)
  58           | Empty -> assert false)
  59     | Empty -> assert false
  60   else Node (l, k, v, r, max hl hr + 1)
  61
  62 let rec min_binding = function
  63   | Node (Empty, k, v, _, _) -> k, v
  64   | Node (l, _, _, _, _) -> min_binding l
  65   | Empty -> raise Not_found
  66
  67 let rec remove_min_binding = function
  68   | Node (Empty, _, _, r, _) -> r
  69   | Node (l, k, v, r, _) -> bal (remove_min_binding l) k v r
  70   | Empty -> invalid_arg "PMap.remove_min_binding"
  71
  72 let merge t1 t2 =
  73   match t1, t2 with
  74   | Empty, _ -> t2
  75   | _, Empty -> t1
  76   | _ ->
  77       let k, v = min_binding t2 in
  78       bal t1 k v (remove_min_binding t2)
  79
  80 let create cmp = { cmp = cmp; map = Empty }
  81 let empty = { cmp = compare; map = Empty }
  82
  83 let is_empty x =
  84         x.map = Empty
  85
  86 let add x d { cmp = cmp; map = map } =
  87   let rec loop = function
  88     | Node (l, k, v, r, h) ->
  89         let c = cmp x k in
  90         if c = 0 then Node (l, x, d, r, h)
  91         else if c < 0 then
  92           let nl = loop l in
  93           bal nl k v r
  94         else
  95           let nr = loop r in
  96           bal l k v nr
  97     | Empty -> Node (Empty, x, d, Empty, 1) in
  98   { cmp = cmp; map = loop map }
  99
 100 let find x { cmp = cmp; map = map } =
 101   let rec loop = function
 102     | Node (l, k, v, r, _) ->
 103         let c = cmp x k in
 104         if c < 0 then loop l
 105         else if c > 0 then loop r
 106         else v
 107     | Empty -> raise Not_found in
 108   loop map
 109
 110 let remove x { cmp = cmp; map = map } =
 111   let rec loop = function
 112     | Node (l, k, v, r, _) ->
 113         let c = cmp x k in
 114         if c = 0 then merge l r else
 115         if c < 0 then bal (loop l) k v r else bal l k v (loop r)
 116     | Empty -> Empty in
 117   { cmp = cmp; map = loop map }
 118
 119 let mem x { cmp = cmp; map = map } =
 120   let rec loop = function
 121     | Node (l, k, v, r, _) ->
 122         let c = cmp x k in
 123         c = 0 || loop (if c < 0 then l else r)
 124     | Empty -> false in
 125   loop map
 126
 127 let exists = mem
 128
 129 let iter f { map = map } =
 130   let rec loop = function
 131     | Empty -> ()
 132     | Node (l, k, v, r, _) -> loop l; f k v; loop r in
 133   loop map
 134
 135 let map f { cmp = cmp; map = map } =
 136   let rec loop = function
 137     | Empty -> Empty
 138     | Node (l, k, v, r, h) ->
 139           let l = loop l in
 140           let r = loop r in
 141           Node (l, k, f v, r, h) in
 142   { cmp = cmp; map = loop map }
 143
 144 let mapi f { cmp = cmp; map = map } =
 145   let rec loop = function
 146     | Empty -> Empty
 147     | Node (l, k, v, r, h) ->
 148           let l = loop l in
 149           let r = loop r in
 150           Node (l, k, f k v, r, h) in
 151   { cmp = cmp; map = loop map }
 152
 153 let fold f { cmp = cmp; map = map } acc =
 154   let rec loop acc = function
 155     | Empty -> acc
 156     | Node (l, k, v, r, _) ->
 157           loop (f v (loop acc l)) r in
 158   loop acc map
 159
 160 let foldi f { cmp = cmp; map = map } acc =
 161   let rec loop acc = function
 162     | Empty -> acc
 163         | Node (l, k, v, r, _) ->
 164        loop (f k v (loop acc l)) r in
 165   loop acc map
 166
 167 let rec enum m =
 168   let rec make l =
 169     let l = ref l in
 170     let rec next() =
 171       match !l with
 172       | [] -> raise Enum.No_more_elements
 173       | Empty :: tl -> l := tl; next()
 174       | Node (m1, key, data, m2, h) :: tl ->
 175         l := m1 :: m2 :: tl;
 176         (key, data)
 177     in
 178     let count() =
 179       let n = ref 0 in
 180       let r = !l in
 181       try
 182         while true do
 183           ignore (next());
 184           incr n
 185         done;
 186         assert false
 187       with
 188                 Enum.No_more_elements -> l := r; !n
 189     in
 190     let clone() = make !l in
 191         Enum.make ~next ~count ~clone
 192   in
 193   make [m.map]
 194
 195
 196 let uncurry_add (k, v) m = add k v m
 197 let of_enum ?(cmp = compare) e = Enum.fold uncurry_add (create cmp) e