commons/ocamlextra/enum.ml

   1 (*
   2  * Enum - Enumeration over abstract collection of elements.
   3  * Copyright (C) 2003 Nicolas Cannasse
   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 'a t = {
  22         mutable count : unit -> int;
  23         mutable next : unit -> 'a;
  24         mutable clone : unit -> 'a t;
  25         mutable fast : bool;
  26 }
  27
  28 (* raised by 'next' functions, should NOT go outside the API *)
  29 exception No_more_elements
  30
  31 let _dummy () = assert false
  32
  33 let make ~next ~count ~clone =
  34         {
  35                 count = count;
  36                 next = next;
  37                 clone = clone;
  38                 fast = true;
  39         }
  40
  41 let rec init n f =
  42         if n < 0 then invalid_arg "Enum.init";
  43         let count = ref n in
  44         {
  45                 count = (fun () -> !count);
  46                 next = (fun () ->
  47                         match !count with
  48                         | 0 -> raise No_more_elements
  49                         | _ ->
  50                                 decr count;
  51                                 f (n - 1 - !count));
  52                 clone = (fun () -> init !count f);
  53                 fast = true;
  54         }
  55
  56 let rec empty () =
  57         {
  58                 count = (fun () -> 0);
  59                 next = (fun () -> raise No_more_elements);
  60                 clone = (fun () -> empty());
  61                 fast = true;
  62         }
  63
  64 type 'a _mut_list = {
  65         hd : 'a;
  66         mutable tl : 'a _mut_list;
  67 }
  68
  69 let force t =
  70         let rec clone enum count =
  71                 let enum = ref !enum
  72                 and     count = ref !count in
  73                 {
  74                         count = (fun () -> !count);
  75                         next = (fun () ->
  76                                 match !enum with
  77                                 | [] -> raise No_more_elements
  78                                 | h :: t -> decr count; enum := t; h);
  79                         clone = (fun () ->
  80                                 let enum = ref !enum
  81                                 and count = ref !count in
  82                                 clone enum count);
  83                         fast = true;
  84                 }
  85         in
  86         let count = ref 0 in
  87         let _empty = Obj.magic [] in
  88         let rec loop dst =
  89                 let x = { hd = t.next(); tl = _empty } in
  90                 incr count;
  91                 dst.tl <- x;
  92                 loop x
  93         in
  94         let enum = ref _empty  in
  95         (try
  96                 enum := { hd = t.next(); tl = _empty };
  97                 incr count;
  98                 loop !enum;
  99         with No_more_elements -> ());
 100         let tc = clone (Obj.magic enum) count in
 101         t.clone <- tc.clone;
 102         t.next <- tc.next;
 103         t.count <- tc.count;
 104         t.fast <- true
 105
 106 let from f =
 107         let e = {
 108                 next = f;
 109                 count = _dummy;
 110                 clone = _dummy;
 111                 fast = false;
 112         } in
 113         e.count <- (fun () -> force e; e.count());
 114         e.clone <- (fun () -> force e; e.clone());
 115         e
 116
 117 let from2 next clone =
 118         let e = {
 119                 next = next;
 120                 count = _dummy;
 121                 clone = clone;
 122                 fast = false;
 123         } in
 124         e.count <- (fun () -> force e; e.count());
 125         e
 126
 127 let get t =
 128         try
 129                 Some (t.next())
 130         with
 131                 No_more_elements -> None
 132
 133 let push t e =
 134         let rec make t =
 135                 let fnext = t.next in
 136                 let fcount = t.count in
 137                 let fclone = t.clone in
 138                 let next_called = ref false in
 139                 t.next <- (fun () ->
 140                         next_called := true;
 141                         t.next <- fnext;
 142                         t.count <- fcount;
 143                         t.clone <- fclone;
 144                         e);
 145                 t.count <- (fun () ->
 146                         let n = fcount() in
 147                         if !next_called then n else n+1);
 148                 t.clone <- (fun () ->
 149                         let tc = fclone() in
 150                         if not !next_called then make tc;
 151                         tc);
 152         in
 153         make t
 154
 155 let peek t =
 156         match get t with
 157         | None -> None
 158         | Some x ->
 159                 push t x;
 160                 Some x
 161
 162 let junk t =
 163         try
 164                 ignore(t.next())
 165         with
 166                 No_more_elements -> ()
 167
 168 let is_empty t =
 169         if t.fast then
 170                 t.count() = 0
 171         else
 172                 peek t = None
 173
 174 let count t =
 175         t.count()
 176
 177 let fast_count t =
 178         t.fast
 179
 180 let clone t =
 181         t.clone()
 182
 183 let iter f t =
 184         let rec loop () =
 185                 f (t.next());
 186                 loop();
 187         in
 188         try
 189                 loop();
 190         with
 191                 No_more_elements -> ()
 192
 193 let iteri f t =
 194         let rec loop idx =
 195                 f idx (t.next());
 196                 loop (idx+1);
 197         in
 198         try
 199                 loop 0;
 200         with
 201                 No_more_elements -> ()
 202
 203 let iter2 f t u =
 204         let push_t = ref None in
 205         let rec loop () =
 206                 push_t := None;
 207                 let e = t.next() in
 208                 push_t := Some e;
 209                 f e (u.next());
 210                 loop ()
 211         in
 212         try
 213                 loop ()
 214         with
 215                 No_more_elements ->
 216                         match !push_t with
 217                         | None -> ()
 218                         | Some e ->
 219                                 push t e
 220
 221 let iter2i f t u =
 222         let push_t = ref None in
 223         let rec loop idx =
 224                 push_t := None;
 225                 let e = t.next() in
 226                 push_t := Some e;
 227                 f idx e (u.next());
 228                 loop (idx + 1)
 229         in
 230         try
 231                 loop 0
 232         with
 233                 No_more_elements ->
 234                         match !push_t with
 235                         | None -> ()
 236                         | Some e -> push t e
 237
 238 let fold f init t =
 239         let acc = ref init in
 240         let rec loop() =
 241                 acc := f (t.next()) !acc;
 242                 loop()
 243         in
 244         try
 245                 loop()
 246         with
 247                 No_more_elements -> !acc
 248
 249 let foldi f init t =
 250         let acc = ref init in
 251         let rec loop idx =
 252                 acc := f idx (t.next()) !acc;
 253                 loop (idx + 1)
 254         in
 255         try
 256                 loop 0
 257         with
 258                 No_more_elements -> !acc
 259
 260 let fold2 f init t u =
 261         let acc = ref init in
 262         let push_t = ref None in
 263         let rec loop() =
 264                 push_t := None;
 265                 let e = t.next() in
 266                 push_t := Some e;
 267                 acc := f e (u.next()) !acc;
 268                 loop()
 269         in
 270         try
 271                 loop()
 272         with
 273                 No_more_elements ->
 274                         match !push_t with
 275                         | None -> !acc
 276                         | Some e ->
 277                                 push t e;
 278                                 !acc
 279
 280 let fold2i f init t u =
 281         let acc = ref init in
 282         let push_t = ref None in
 283         let rec loop idx =
 284                 push_t := None;
 285                 let e = t.next() in
 286                 push_t := Some e;
 287                 acc := f idx e (u.next()) !acc;
 288                 loop (idx + 1)
 289         in
 290         try
 291                 loop 0
 292         with
 293                 No_more_elements ->
 294                         match !push_t with
 295                         | None -> !acc
 296                         | Some e ->
 297                                 push t e;
 298                                 !acc
 299
 300 let find f t =
 301         let rec loop () =
 302                 let x = t.next() in
 303                 if f x then x else loop()
 304         in
 305         try
 306                 loop()
 307         with
 308                 No_more_elements -> raise Not_found
 309
 310 let rec map f t =
 311         {
 312                 count = t.count;
 313                 next = (fun () -> f (t.next()));
 314                 clone = (fun () -> map f (t.clone()));
 315                 fast = t.fast;
 316         }
 317
 318 let rec mapi f t =
 319         let idx = ref (-1) in
 320         {
 321                 count = t.count;
 322                 next = (fun () -> incr idx; f !idx (t.next()));
 323                 clone = (fun () -> mapi f (t.clone()));
 324                 fast = t.fast;
 325         }
 326
 327 let rec filter f t =
 328         let rec next() =
 329                 let x = t.next() in
 330                 if f x then x else next()
 331         in
 332         from2 next (fun () -> filter f (t.clone()))
 333
 334 let rec filter_map f t =
 335     let rec next () =
 336         match f (t.next()) with
 337         | None -> next()
 338         | Some x -> x
 339     in
 340         from2 next (fun () -> filter_map f (t.clone()))
 341
 342 let rec append ta tb =
 343         let t = {
 344                 count = (fun () -> ta.count() + tb.count());
 345                 next = _dummy;
 346                 clone = (fun () -> append (ta.clone()) (tb.clone()));
 347                 fast = ta.fast && tb.fast;
 348         } in
 349         t.next <- (fun () ->
 350                 try
 351                         ta.next()
 352                 with
 353                         No_more_elements ->
 354                                 (* add one indirection because tb can mute *)
 355                                 t.next <- (fun () -> tb.next());
 356                                 t.count <- (fun () -> tb.count());
 357                                 t.clone <- (fun () -> tb.clone());
 358                                 t.fast <- tb.fast;
 359                                 t.next()
 360         );
 361         t
 362
 363 let rec concat t =
 364         let concat_ref = ref _dummy in
 365         let rec concat_next() =
 366                 let tn = t.next() in
 367                 concat_ref := (fun () ->
 368                         try
 369                                 tn.next()
 370                         with
 371                                 No_more_elements ->
 372                                         concat_next());
 373                 !concat_ref ()
 374         in
 375         concat_ref := concat_next;
 376         from2 (fun () -> !concat_ref ()) (fun () -> concat (t.clone()))