760 lines
16 KiB
OCaml
760 lines
16 KiB
OCaml
(* Lightweight thread library for Objective Caml
|
|
* http://www.ocsigen.org/lwt
|
|
* Module Lwt_stream
|
|
* Copyright (C) 2009 Jérémie Dimino
|
|
*
|
|
* This program is free software; you can redistribute it and/or modify
|
|
* it under the terms of the GNU Lesser General Public License as
|
|
* published by the Free Software Foundation, with linking exceptions;
|
|
* either version 2.1 of the License, or (at your option) any later
|
|
* version. See COPYING file for details.
|
|
*
|
|
* This program is distributed in the hope that it will be useful, but
|
|
* WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
|
|
* Lesser General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU Lesser General Public
|
|
* License along with this program; if not, write to the Free Software
|
|
* Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
|
|
* 02111-1307, USA.
|
|
*)
|
|
|
|
open Lwt
|
|
|
|
exception Empty
|
|
|
|
type 'a t = {
|
|
next : unit -> 'a option Lwt.t;
|
|
(* The source of the stream *)
|
|
queue : 'a option Queue.t;
|
|
(* Queue of pending elements, which are not yet consumed *)
|
|
clones : 'a option Queue.t Weak.t ref;
|
|
(* List of queues of all clones of this event (including this
|
|
event) *)
|
|
mutex : Lwt_mutex.t;
|
|
(* Mutex to prevent concurrent access to [next] *)
|
|
}
|
|
|
|
let add_clone wa q =
|
|
let len = Weak.length !wa in
|
|
(* loop search for a free cell in [wa] and fill it with [q]: *)
|
|
let rec loop i =
|
|
if i = len then begin
|
|
(* Growing *)
|
|
let clones = Weak.create (len + 1) in
|
|
Weak.blit !wa 0 clones 0 len;
|
|
wa := clones;
|
|
Weak.set clones len (Some q)
|
|
end else if Weak.check !wa i then
|
|
loop (i + 1)
|
|
else
|
|
Weak.set !wa i (Some q)
|
|
in
|
|
loop 0
|
|
|
|
let clone s =
|
|
let s' = {
|
|
next = s.next;
|
|
queue = Queue.copy s.queue;
|
|
clones = s.clones;
|
|
mutex = s.mutex;
|
|
} in
|
|
add_clone s'.clones s'.queue;
|
|
s'
|
|
|
|
let from f =
|
|
let s = {
|
|
next = f;
|
|
queue = Queue.create ();
|
|
clones = ref(Weak.create 1);
|
|
mutex = Lwt_mutex.create ();
|
|
} in
|
|
Weak.set !(s.clones) 0 (Some s.queue);
|
|
s
|
|
|
|
let of_list l =
|
|
let l = ref l in
|
|
from (fun () ->
|
|
match !l with
|
|
| [] -> return None
|
|
| x :: l' -> l := l'; return (Some x))
|
|
|
|
let of_array a =
|
|
let len = Array.length a and i = ref 0 in
|
|
from (fun () ->
|
|
if !i = len then
|
|
return None
|
|
else begin
|
|
let c = Array.unsafe_get a !i in
|
|
incr i;
|
|
return (Some c)
|
|
end)
|
|
|
|
let of_string s =
|
|
let len = String.length s and i = ref 0 in
|
|
from (fun () ->
|
|
if !i = len then
|
|
return None
|
|
else begin
|
|
let c = String.unsafe_get s !i in
|
|
incr i;
|
|
return (Some c)
|
|
end)
|
|
|
|
module EQueue :
|
|
sig
|
|
type 'a t
|
|
val create : unit -> 'a t * ('a option -> unit)
|
|
val pop : 'a t -> 'a option Lwt.t
|
|
end =
|
|
struct
|
|
type 'a state =
|
|
| No_mail
|
|
| Waiting of 'a option Lwt.u
|
|
| Full of 'a option Queue.t
|
|
|
|
type 'a t = {
|
|
mutable state : 'a state;
|
|
}
|
|
|
|
let create () =
|
|
let box = { state = No_mail } in
|
|
let weak_box = Weak.create 1 in
|
|
Weak.set weak_box 0 (Some box);
|
|
let push v =
|
|
match Weak.get weak_box 0 with
|
|
| None -> ()
|
|
| Some box ->
|
|
match box.state with
|
|
| No_mail ->
|
|
let q = Queue.create () in
|
|
Queue.push v q;
|
|
box.state <- Full q
|
|
| Waiting wakener ->
|
|
box.state <- No_mail;
|
|
wakeup_later wakener v
|
|
| Full q ->
|
|
Queue.push v q
|
|
in
|
|
(box, push)
|
|
|
|
let pop b = match b.state with
|
|
| No_mail ->
|
|
let waiter, wakener = task () in
|
|
Lwt.on_cancel waiter (fun () -> b.state <- No_mail);
|
|
b.state <- Waiting wakener;
|
|
waiter
|
|
| Waiting _ ->
|
|
(* Calls to next are serialized, so this case will never
|
|
happened *)
|
|
assert false
|
|
| Full q ->
|
|
let v = Queue.take q in
|
|
if Queue.is_empty q then b.state <- No_mail;
|
|
return v
|
|
end
|
|
|
|
let create () =
|
|
let box, push = EQueue.create () in
|
|
(from (fun () -> EQueue.pop box), push)
|
|
|
|
let push_clones wa x =
|
|
for i = 0 to Weak.length wa - 1 do
|
|
match Weak.get wa i with
|
|
| Some q ->
|
|
Queue.push x q
|
|
| None ->
|
|
()
|
|
done
|
|
|
|
let peek s =
|
|
if Queue.is_empty s.queue then
|
|
Lwt_mutex.with_lock s.mutex begin fun () ->
|
|
if Queue.is_empty s.queue then begin
|
|
lwt result = s.next () in
|
|
push_clones !(s.clones) result;
|
|
return result
|
|
end else
|
|
return (Queue.top s.queue)
|
|
end
|
|
else
|
|
return (Queue.top s.queue)
|
|
|
|
let rec force n s =
|
|
if Queue.length s.queue >= n then
|
|
return ()
|
|
else
|
|
Lwt_mutex.with_lock s.mutex begin fun () ->
|
|
if Queue.length s.queue >= n then
|
|
return false
|
|
else begin
|
|
lwt result = s.next () in
|
|
push_clones !(s.clones) result;
|
|
if result = None then
|
|
return false
|
|
else
|
|
return true
|
|
end
|
|
end >>= function
|
|
| true ->
|
|
force n s
|
|
| false ->
|
|
return ()
|
|
|
|
let npeek n s =
|
|
lwt () = force n s in
|
|
let l, _ =
|
|
Queue.fold
|
|
(fun (l, n) x ->
|
|
if n > 0 then
|
|
match x with
|
|
| Some x -> (x :: l, n - 1)
|
|
| None -> (l, n)
|
|
else
|
|
(l, n))
|
|
([], n) s.queue
|
|
in
|
|
return (List.rev l)
|
|
|
|
let rec get s =
|
|
if Queue.is_empty s.queue then
|
|
Lwt_mutex.with_lock s.mutex begin fun () ->
|
|
if Queue.is_empty s.queue then begin
|
|
lwt x = s.next () in
|
|
(* This prevent from calling s.next when the stream has
|
|
terminated: *)
|
|
if x = None then Queue.push None s.queue;
|
|
let wa = !(s.clones) in
|
|
for i = 0 to Weak.length wa - 1 do
|
|
match Weak.get wa i with
|
|
| Some q when q != s.queue ->
|
|
Queue.push x q
|
|
| _ ->
|
|
()
|
|
done;
|
|
return x
|
|
end else begin
|
|
let x = Queue.take s.queue in
|
|
if x = None then Queue.push None s.queue;
|
|
return x
|
|
end
|
|
end
|
|
else begin
|
|
let x = Queue.take s.queue in
|
|
if x = None then Queue.push None s.queue;
|
|
return x
|
|
end
|
|
|
|
let nget n s =
|
|
let rec loop n =
|
|
if n <= 0 then
|
|
return []
|
|
else
|
|
get s >>= function
|
|
| Some x ->
|
|
lwt l = loop (n - 1) in
|
|
return (x :: l)
|
|
| None ->
|
|
return []
|
|
in
|
|
loop n
|
|
|
|
let get_while f s =
|
|
let rec loop () =
|
|
peek s >>= function
|
|
| Some x ->
|
|
let test = f x in
|
|
if test then begin
|
|
ignore (Queue.take s.queue);
|
|
lwt l = loop () in
|
|
return (x :: l)
|
|
end else
|
|
return []
|
|
| None ->
|
|
return []
|
|
in
|
|
loop ()
|
|
|
|
let get_while_s f s =
|
|
let rec loop () =
|
|
peek s >>= function
|
|
| Some x ->
|
|
lwt test = f x in
|
|
if test then begin
|
|
ignore (Queue.take s.queue);
|
|
lwt l = loop () in
|
|
return (x :: l)
|
|
end else
|
|
return []
|
|
| None ->
|
|
return []
|
|
in
|
|
loop ()
|
|
|
|
let next s = get s >>= function
|
|
| Some x -> return x
|
|
| None -> raise_lwt Empty
|
|
|
|
let last_new s =
|
|
match Lwt.state (peek s) with
|
|
| Return None ->
|
|
raise_lwt Empty
|
|
| Sleep ->
|
|
next s
|
|
| Fail exn ->
|
|
raise_lwt exn
|
|
| Return(Some x) ->
|
|
let _ = Queue.take s.queue in
|
|
let rec loop last =
|
|
match Lwt.state (peek s) with
|
|
| Sleep | Return None ->
|
|
return last
|
|
| Return(Some x) ->
|
|
let _ = Queue.take s.queue in
|
|
loop x
|
|
| Fail exn ->
|
|
raise_lwt exn
|
|
in
|
|
loop x
|
|
|
|
let to_list s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
lwt l = loop () in
|
|
return (x :: l)
|
|
| None ->
|
|
return []
|
|
in
|
|
loop ()
|
|
|
|
let to_string s =
|
|
let buf = Buffer.create 42 in
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
Buffer.add_char buf x;
|
|
loop ()
|
|
| None ->
|
|
return (Buffer.contents buf)
|
|
in
|
|
loop ()
|
|
|
|
let junk s =
|
|
lwt _ = get s in
|
|
return ()
|
|
|
|
let njunk n s =
|
|
let rec loop n =
|
|
if n <= 0 then
|
|
return ()
|
|
else
|
|
lwt _ = get s in
|
|
loop (n - 1)
|
|
in
|
|
loop n
|
|
|
|
let junk_while f s =
|
|
let rec loop () =
|
|
peek s >>= function
|
|
| Some x ->
|
|
let test = f x in
|
|
if test then begin
|
|
ignore (Queue.take s.queue);
|
|
loop ()
|
|
end else
|
|
return ()
|
|
| None ->
|
|
return ()
|
|
in
|
|
loop ()
|
|
|
|
let junk_while_s f s =
|
|
let rec loop () =
|
|
peek s >>= function
|
|
| Some x ->
|
|
lwt test = f x in
|
|
if test then begin
|
|
ignore (Queue.take s.queue);
|
|
loop ()
|
|
end else
|
|
return ()
|
|
| None ->
|
|
return ()
|
|
in
|
|
loop ()
|
|
|
|
let junk_old s =
|
|
let rec loop () =
|
|
match Lwt.state (peek s) with
|
|
| Sleep ->
|
|
return ()
|
|
| _ ->
|
|
ignore (Queue.take s.queue);
|
|
loop ()
|
|
in
|
|
loop ()
|
|
|
|
let get_available s =
|
|
let rec loop () =
|
|
match Lwt.state (peek s) with
|
|
| Sleep | Return None ->
|
|
[]
|
|
| Return(Some x) ->
|
|
ignore (Queue.take s.queue);
|
|
x :: loop ()
|
|
| Fail exn ->
|
|
raise exn
|
|
in
|
|
loop ()
|
|
|
|
let get_available_up_to n s =
|
|
let rec loop = function
|
|
| 0 ->
|
|
[]
|
|
| n ->
|
|
match Lwt.state (peek s) with
|
|
| Sleep | Return None ->
|
|
[]
|
|
| Return(Some x) ->
|
|
ignore (Queue.take s.queue);
|
|
x :: loop (n - 1)
|
|
| Fail exn ->
|
|
raise exn
|
|
in
|
|
loop n
|
|
|
|
let is_empty s = peek s >|= fun x -> x = None
|
|
|
|
let map f s =
|
|
from (fun () -> get s >>= function
|
|
| Some x ->
|
|
let x = f x in
|
|
return (Some x)
|
|
| None ->
|
|
return None)
|
|
|
|
let map_s f s =
|
|
from (fun () -> get s >>= function
|
|
| Some x ->
|
|
lwt x = f x in
|
|
return (Some x)
|
|
| None ->
|
|
return None)
|
|
|
|
let filter f s =
|
|
let rec next () =
|
|
get s >>= function
|
|
| Some x as result ->
|
|
let test = f x in
|
|
if test then
|
|
return result
|
|
else
|
|
next ()
|
|
| None ->
|
|
return None
|
|
in
|
|
from next
|
|
|
|
let filter_s f s =
|
|
let rec next () =
|
|
get s >>= function
|
|
| Some x as result ->
|
|
lwt test = f x in
|
|
if test then
|
|
return result
|
|
else
|
|
next ()
|
|
| None ->
|
|
return None
|
|
in
|
|
from next
|
|
|
|
let filter_map f s =
|
|
let rec next () =
|
|
get s >>= function
|
|
| Some x ->
|
|
let x = f x in
|
|
(match x with
|
|
| Some _ ->
|
|
return x
|
|
| None ->
|
|
next ())
|
|
| None ->
|
|
return None
|
|
in
|
|
from next
|
|
|
|
let filter_map_s f s =
|
|
let rec next () =
|
|
get s >>= function
|
|
| Some x ->
|
|
lwt x = f x in
|
|
(match x with
|
|
| Some _ ->
|
|
return x
|
|
| None ->
|
|
next ())
|
|
| None ->
|
|
return None
|
|
in
|
|
from next
|
|
|
|
let map_list f s =
|
|
let pendings = ref [] in
|
|
let rec next () =
|
|
match !pendings with
|
|
| [] ->
|
|
get s >>= (function
|
|
| Some x ->
|
|
let l = f x in
|
|
pendings := l;
|
|
next ()
|
|
| None ->
|
|
return None)
|
|
| x :: l ->
|
|
pendings := l;
|
|
return (Some x)
|
|
in
|
|
from next
|
|
|
|
let map_list_s f s =
|
|
let pendings = ref [] in
|
|
let rec next () =
|
|
match !pendings with
|
|
| [] ->
|
|
get s >>= (function
|
|
| Some x ->
|
|
lwt l = f x in
|
|
pendings := l;
|
|
next ()
|
|
| None ->
|
|
return None)
|
|
| x :: l ->
|
|
pendings := l;
|
|
return (Some x)
|
|
in
|
|
from next
|
|
|
|
let flatten s =
|
|
map_list (fun l -> l) s
|
|
|
|
let fold f s acc =
|
|
let rec loop acc =
|
|
get s >>= function
|
|
| Some x ->
|
|
let acc = f x acc in
|
|
loop acc
|
|
| None ->
|
|
return acc
|
|
in
|
|
loop acc
|
|
|
|
let fold_s f s acc =
|
|
let rec loop acc =
|
|
get s >>= function
|
|
| Some x ->
|
|
lwt acc = f x acc in
|
|
loop acc
|
|
| None ->
|
|
return acc
|
|
in
|
|
loop acc
|
|
|
|
let iter f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
let () = f x in
|
|
loop ()
|
|
| None ->
|
|
return ()
|
|
in
|
|
loop ()
|
|
|
|
let iter_s f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
lwt () = f x in
|
|
loop ()
|
|
| None ->
|
|
return ()
|
|
in
|
|
loop ()
|
|
|
|
let iter_p f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
f x <&> loop ()
|
|
| None ->
|
|
return ()
|
|
in
|
|
loop ()
|
|
|
|
let find f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x as result ->
|
|
let test = f x in
|
|
if test then
|
|
return result
|
|
else
|
|
loop ()
|
|
| None ->
|
|
return None
|
|
in
|
|
loop ()
|
|
|
|
let find_s f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x as result ->
|
|
lwt test = f x in
|
|
if test then
|
|
return result
|
|
else
|
|
loop ()
|
|
| None ->
|
|
return None
|
|
in
|
|
loop ()
|
|
|
|
let rec find_map f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
let x = f x in
|
|
(match x with
|
|
| Some _ ->
|
|
return x
|
|
| None ->
|
|
loop ())
|
|
| None ->
|
|
return None
|
|
in
|
|
loop ()
|
|
|
|
let rec find_map_s f s =
|
|
let rec loop () =
|
|
get s >>= function
|
|
| Some x ->
|
|
lwt x = f x in
|
|
(match x with
|
|
| Some _ ->
|
|
return x
|
|
| None ->
|
|
loop ())
|
|
| None ->
|
|
return None
|
|
in
|
|
loop ()
|
|
|
|
let rec combine s1 s2 =
|
|
let next () =
|
|
lwt n1 = get s1 and n2 = get s2 in
|
|
match n1, n2 with
|
|
| Some x1, Some x2 ->
|
|
return (Some(x1, x2))
|
|
| _ ->
|
|
return None
|
|
in
|
|
from next
|
|
|
|
let append s1 s2 =
|
|
let current_s = ref s1 and s1_finished = ref false in
|
|
let rec next () =
|
|
get !current_s >>= function
|
|
| Some _ as result ->
|
|
return result
|
|
| None ->
|
|
if !s1_finished then
|
|
return None
|
|
else begin
|
|
s1_finished := true;
|
|
current_s := s2;
|
|
next ()
|
|
end
|
|
in
|
|
from next
|
|
|
|
let concat s_top =
|
|
let current_s = ref(from(fun () -> return None)) in
|
|
let rec next () =
|
|
get !current_s >>= function
|
|
| Some _ as result ->
|
|
return result
|
|
| None ->
|
|
get s_top >>= function
|
|
| Some s ->
|
|
current_s := s;
|
|
next ()
|
|
| None ->
|
|
return None
|
|
in
|
|
from next
|
|
|
|
let choose streams =
|
|
let source s = (s, peek s >|= fun x -> (s, x)) in
|
|
let streams = ref (List.rev_map source streams) in
|
|
let rec next () =
|
|
match !streams with
|
|
| [] ->
|
|
return None
|
|
| l ->
|
|
lwt s, x = Lwt.choose (List.map snd l) in
|
|
let l = List.remove_assq s l in
|
|
match x with
|
|
| Some _ ->
|
|
lwt () = junk s in
|
|
streams := source s :: l;
|
|
return x
|
|
| None ->
|
|
next ()
|
|
in
|
|
from next
|
|
|
|
let parse s f =
|
|
let s' = clone s in
|
|
try_lwt
|
|
f s
|
|
with exn ->
|
|
Queue.clear s.queue;
|
|
Queue.transfer s'.queue s.queue;
|
|
raise_lwt exn
|
|
|
|
let hexdump stream =
|
|
let buf = Buffer.create 80 and num = ref 0 in
|
|
from begin fun _ ->
|
|
nget 16 stream >>= function
|
|
| [] ->
|
|
return None
|
|
| l ->
|
|
Buffer.clear buf;
|
|
Printf.bprintf buf "%08x| " !num;
|
|
num := !num + 16;
|
|
let rec bytes pos = function
|
|
| [] ->
|
|
blanks pos
|
|
| x :: l ->
|
|
if pos = 8 then Buffer.add_char buf ' ';
|
|
Printf.bprintf buf "%02x " (Char.code x);
|
|
bytes (pos + 1) l
|
|
and blanks pos =
|
|
if pos < 16 then begin
|
|
if pos = 8 then
|
|
Buffer.add_string buf " "
|
|
else
|
|
Buffer.add_string buf " ";
|
|
blanks (pos + 1)
|
|
end
|
|
in
|
|
bytes 0 l;
|
|
Buffer.add_string buf " |";
|
|
List.iter (fun ch -> Buffer.add_char buf (if ch >= '\x20' && ch <= '\x7e' then ch else '.')) l;
|
|
Buffer.add_char buf '|';
|
|
return (Some(Buffer.contents buf))
|
|
end
|