1 files changed, 51 insertions, 41 deletions

diff --git a/stdlib/stream.ml b/stdlib/stream.ml
index c9975c4a0..aa6a2a2ab 100644
--- a/stdlib/stream.ml
+++ b/stdlib/stream.ml
@@ -22,7 +22,7 @@ and 'a data =
     Sempty
   | Scons of 'a * 'a data
   | Sapp of 'a data * 'a data
-  | Slazy of (unit -> 'a data)
+  | Slazy of 'a data Lazy.t
   | Sgen of 'a gen
   | Sbuffio of buffio
 and 'a gen = { mutable curr : 'a option option; func : int -> 'a option }
@@ -42,44 +42,54 @@ let fill_buff b =
   b.len <- input b.ic b.buff 0 (String.length b.buff); b.ind <- 0
 ;;
 
-let rec get_data =
-  function
-    Sempty -> None
-  | Scons (a, d) -> Some (a, d)
-  | Sapp (d1, d2) ->
-      begin match get_data d1 with
-        Some (a, d1) -> Some (a, Sapp (d1, d2))
-      | None -> get_data d2
-      end
-  | Slazy f ->
-      begin match f () with
-        Sgen _ | Sbuffio _ -> failwith "illegal stream concatenation"
-      | x -> get_data x
-      end
-  | Sgen _ | Sbuffio _ ->
-      failwith "illegal stream concatenation"
+let rec get_data count d = match d with
+ (* Returns either Sempty or Scons(a, _) even when d is a generator
+    or a buffer. In those cases, the item a is seen as extracted from
+ the generator/buffer.
+ The count parameter is used for calling `Sgen-functions'.  *)
+   Sempty | Scons (_, _) -> d
+ | Sapp (d1, d2) ->
+     begin match get_data count d1 with
+       Scons (a, d11) -> Scons (a, Sapp (d11, d2))
+     | Sempty -> get_data count d2
+     | _ -> assert false
+     end
+ | Sgen {curr = Some None; func = _ } -> Sempty
+ | Sgen ({curr = Some(Some a); func = f} as g) ->
+     g.curr <- None; Scons(a, d)
+ | Sgen g ->
+     begin match g.func count with
+       None -> g.curr <- Some(None); Sempty
+     | Some a -> Scons(a, d)
+         (* Warning: anyone using g thinks that an item has been read *)
+     end
+ | Sbuffio b ->
+     if b.ind >= b.len then fill_buff b;
+     if b.len == 0 then Sempty else
+       let r = Obj.magic (String.unsafe_get b.buff b.ind) in
+       (* Warning: anyone using g thinks that an item has been read *)
+       b.ind <- succ b.ind; Scons(r, d)
+ | Slazy f -> get_data count (Lazy.force f)
 ;;
 
 let rec peek s =
-  match s.data with
-    Sempty -> None
-  | Scons (a, _) -> Some a
-  | Sapp (_, _) ->
-      begin match get_data s.data with
-        Some (a, d) -> set_data s (Scons (a, d)); Some a
-      | None -> None
-      end
-  | Slazy f ->
-      begin match f () with
-        Sgen _ | Sbuffio _ -> failwith "illegal stream concatenation"
-      | d -> set_data s d; peek s
-      end
-  | Sgen {curr = Some a} -> a
-  | Sgen g -> let x = g.func s.count in g.curr <- Some x; x
-  | Sbuffio b ->
-      if b.ind >= b.len then fill_buff b;
-      if b.len == 0 then begin set_data s Sempty; None end
-      else Some (Obj.magic (String.unsafe_get b.buff b.ind))
+ (* consult the first item of s *)
+ match s.data with
+   Sempty -> None
+ | Scons (a, _) -> Some a
+ | Sapp (_, _) ->
+     begin match get_data s.count s.data with
+       Scons(a, _) as d -> set_data s d; Some a
+     | Sempty -> None
+     | _ -> assert false
+     end
+ | Slazy f -> set_data s (Lazy.force f); peek s
+ | Sgen {curr = Some a} -> a
+ | Sgen g -> let x = g.func s.count in g.curr <- Some x; x
+ | Sbuffio b ->
+     if b.ind >= b.len then fill_buff b;
+     if b.len == 0 then begin set_data s Sempty; None end
+     else Some (Obj.magic (String.unsafe_get b.buff b.ind))
 ;;
 
 let rec junk s =
@@ -152,13 +162,13 @@ let icons i s = {count = 0; data = Scons (i, s.data)};;
 let ising i = {count = 0; data = Scons (i, Sempty)};;
 
 let lapp f s =
-  {count = 0; data = Slazy (fun _ -> Sapp ((f ()).data, s.data))}
+  {count = 0; data = Slazy (lazy(Sapp ((f ()).data, s.data)))}
 ;;
-let lcons f s = {count = 0; data = Slazy (fun _ -> Scons (f (), s.data))};;
-let lsing f = {count = 0; data = Slazy (fun _ -> Scons (f (), Sempty))};;
+let lcons f s = {count = 0; data = Slazy (lazy(Scons (f (), s.data)))};;
+let lsing f = {count = 0; data = Slazy (lazy(Scons (f (), Sempty)))};;
 
 let sempty = {count = 0; data = Sempty};;
-let slazy f = {count = 0; data = Slazy (fun _ -> (f ()).data)};;
+let slazy f = {count = 0; data = Slazy (lazy(f ()).data)};;
 
 (* For debugging use *)
 
@@ -184,7 +194,7 @@ and dump_data f =
       print_string ", ";
       dump_data f d2;
       print_string ")"
-  | Slazy f -> print_string "Slazy"
+  | Slazy _ -> print_string "Slazy"
   | Sgen _ -> print_string "Sgen"
   | Sbuffio b -> print_string "Sbuffio"
 ;;