(***********************************************************************)
(*                                                                     *)
(*                             Ocaml                                   *)
(*                                                                     *)
(*                           Tim Freeman                               *)
(*                                                                     *)
(*  Copyright 2002 Institut National de Recherche en Informatique et   *)
(*  en Automatique.  All rights reserved.  This file is distributed    *)
(*  under the terms of the GNU Library General Public License, with    *)
(*  the special exception on linking described in file ../LICENSE.     *)
(*                                                                     *)
(***********************************************************************)

(* $Id: stream.ml,v 1.13 2001/12/07 13:40:59 xleroy Exp $ *)

(* This rewrite is tracked as bug 1284 at http://caml.inria.fr/bin/caml-bugs.
   Bug 235 is also mentioned. *) 

type 'a body = 
    {peek: unit -> 'a option;
     (* junk discards the first element from the stream.  For efficiency, it
        may set the "body" field of the given wrap to a new value. *)
     junk: 'a wrap -> unit;
     (* printtype is just used for dumping. *)
     printtype: string}

and 'a wrap =
    {mutable count: int;
     (* Support modifying it in place so we can implement backing up a stream
        in place for npeek. *) 
     mutable body: 'a body}

(* Special case for the empty stream so it can be polymorphic. *)
type 'a t = Empty | Nonempty of 'a wrap
                
exception Failure

exception Error of string

let wrap (b: 'a body): 'a t = Nonempty {count = 0; body = b}

let from (f: int -> 'a option): 'a t =
  let pos: int ref = ref 0
  in let current: 'a option option ref = ref None
  in let result: 'a body =
      {peek = (function _ ->
                 match !current with
                     None ->
                       let peeked = f !pos in
                         current := Some peeked;
                         peeked
                   | Some s -> s);
       junk = (function _ ->
                 current := None;
                 incr pos);
       printtype = "from"}
  in wrap result

let of_list (l: 'a list): 'a t =
  let l = ref l in
  let c = ref 0 in
  let result: 'a body =
    {peek = (function _ ->
               match !l with
                   [] -> None
                 | hd :: tail -> Some hd);
     junk = (function _ -> begin
               match !l with
                   [] -> ()
                 | hd :: tail ->
                     incr c;
                     l := tail;
             end);
     printtype = "of_list"}
  in wrap result

let of_string (s: string): char t =
  let pos = ref 0 in
  let result: 'a body =
    {peek = (function _ ->
               if !pos >= String.length s then
                 None
               else
                 Some (s.[!pos]));
     junk = (function _ ->
               incr pos);
     printtype = "of_string"}
  in wrap result

let of_channel (inch: in_channel): char t =
  let pos = ref 0 (* Position in the buffer. *)
  and len = ref 0 (* Number of chars in the buffer. *)
  and eof = ref false
              (* Whether we've closed out inch.  If !eof is true, then
                 len is 0.  For the first byte, !len is 0 and !eof is false. *)
  and buff = String.create 4096
  in let result: 'a body =
      {peek = (function _ ->
                 if !pos < !len then
                   (* String.unsafe_get saves 7% on the timing run over
                      buff.[!pos]. *)
                   (* Some (buff.[!pos]) *)
                   Some (String.unsafe_get buff !pos)
                 else begin
                   while !len <= !pos && not !eof do
                     pos := !pos - !len;
                     len := input inch buff 0 (String.length buff);
                     if 0 = !len then 
                       eof := true;
                   done;
                   if !eof then 
                     None
                   else
                     Some (String.unsafe_get buff !pos)
                 end);
       junk = (function _ ->
                 incr pos);
       printtype = "of_channel"}
  in wrap result

let peek: 'a t -> 'a option = function
  | Empty -> None
  | Nonempty n -> n.body.peek ()
  
let junk: 'a t -> unit = function
  | Empty -> ()
  | Nonempty n ->
      n.count <- n.count + 1;
      n.body.junk n
  
let next (s: 'a t): 'a =
  match peek s with
      Some e ->
        junk s;
        e
    | None -> raise Failure

let empty (s: 'a t): unit =
  match peek s with
      Some s -> raise Failure
    | None -> ()

let count: 'a t -> int = function
    Empty -> 0
  | Nonempty n -> n.count

let rec iter (f: 'a -> unit) (s: 'a t):  unit =
  match peek s with
      None -> ()
    | Some e ->
        f e;
        junk s;
        iter f s

let sempty: 'a t = Empty

let getwrap: 'a t -> 'a wrap = function
    Empty -> {count = 0;
              body = {peek = (function () -> None);
                      junk = (function _ -> ());
                      printtype = "Empty"}}
  | Nonempty n -> n

let iapp (s1: 'a t) (s2: 'a t): 'a t =
  let wrap1: 'a wrap = getwrap s1
  in let wrap2: 'a wrap = getwrap s2
  in let result =
      {peek = (function _ ->
                 match wrap1.body.peek () with
                     Some s -> Some s
                   | None -> wrap2.body.peek ());
       junk = (function wrap ->
                 match wrap1.body.peek() with
                     Some s ->
                       wrap1.body.junk wrap1
                   | None ->
                       wrap.body <- wrap2.body;
                       wrap2.body.junk wrap);
       printtype = "iapp"}
  in wrap result

let icons_body (i: 'a) (b: 'a body): 'a body =
  {peek = (function _ -> Some i);
   junk = (function wrap ->
             wrap.body <- b);
   printtype = "icons"}

let icons (i: 'a) (s: 'a t): 'a t = wrap (icons_body i (getwrap s).body)
    
let ising (i: 'a): 'a t = icons i sempty
                            
let lapp (f: unit -> 'a t) (s2: 'a t): 'a t =
  let fresult: 'a wrap option ref = ref None
  in let getw1 (): 'a wrap =
      match !fresult with
          None ->
            let result = getwrap (f ())
            in
              fresult := Some result;
              result
        | Some w -> w
  in let w2 = getwrap s2
  in let me =
    {peek = (function _ -> 
               match (getw1()).body.peek() with
                   Some x -> Some x
                 | None -> w2.body.peek());
     junk = (function (w: 'a wrap) ->
               let w1: 'a wrap = getw1 () in
                 match w1.body.peek () with
                     Some x ->
                       w1.body.junk w1
                   | None ->
                       w.body <- w2.body;
                       w2.body.junk w);
    printtype = "lapp"}
  in wrap me

let lcons (f: unit -> 'a) (s2: 'a t): 'a t =
  let hd_to_use: 'a option ref = ref None
  in let get_hd () =
      match !hd_to_use with
          None ->
            let result = f () in
              hd_to_use := Some result;
              result
        | Some hd -> hd
  in
    wrap {peek = (function _ -> Some (get_hd()));
          junk = (function w ->
                    w.body <- (getwrap s2).body);
          printtype = "lcons"}
    
let lsing (f: unit -> 'a): 'a t = lcons f sempty

let slazy (f: unit -> 'a t): 'a t =
  let w_to_use: 'a wrap option ref = ref None
  in let get_w () =
      match !w_to_use with
          None ->
            let result = getwrap (f ()) in
              w_to_use := Some result;
              result;
        | Some s -> s
  in
    wrap {peek = (function _ -> ((get_w ()).body.peek ()));
          junk = (function w ->
                    w.body <- (get_w ()).body;
                    w.body.junk w);
          printtype = "slazy"}

(* Like npeek, except it junks the elements it removes instead of putting them
   back, and it starts with a wrap. *)
let rec njunk (i: int) (w: 'a wrap): 'a list =
  if 0 == i then []
  else
    match w.body.peek() with
      Some e ->
        w.body.junk w;
        e :: njunk (i - 1) w
    | None -> []

let rec ncons (l: 'a list) (w: 'a wrap): unit =
  match l with
      [] -> ()
    | hd :: tail ->
        ncons tail w;
        w.body <- icons_body hd w.body

let npeek (i: int) (s: 'a t): 'a list =
  match s with
      Empty -> []
    | Nonempty n ->
        let l = njunk i n
        in
          ncons l n;
          l

let rec dump (dump_data: ('a -> unit)) (s: 'a t): unit =
  print_string "type=";
  print_string (getwrap s).body.printtype;
  print_string " count=";
  print_int (count s);
  print_string " value=";
  begin
    match peek s with
        None ->
          print_string "None";
          print_string "\n";
      | Some e ->
          print_string "Some ";
          dump_data e;
          print_string "\n";
          junk s;
          dump dump_data s;
  end;
  
        
(* Test code: 
   
#load "stream_fixed.cmo";;
  
open Stream_fixed;;

let dump_string s = print_string ("\"" ^ s ^ "\"");;
let mydump = dump dump_string;;
let dump_char c = print_string (Printf.sprintf "%c" c);;
let checkfailure (f: unit -> 'a): unit =
    begin
      try
        ignore (f ());
        raise (Error "Should have thrown")
      with Failure ->
        print_string "Test passed\n";
        ()
    end;;

(* Test cases are separated by blank lines. *)
let big_test () = begin
  
  mydump sempty;
  
  let fromfun (x: int): string option =
    if x < 10 then
      Some (Printf.sprintf "%d" x)
    else
      None
  in begin
      mydump (from fromfun);
      assert (npeek 20 (from fromfun) =
                ["0"; "1"; "2"; "3"; "4"; "5"; "6"; "7"; "8"; "9"]);
    end;

  let lstream () = of_list [1; 2; 3; 4] in
  let dump_int s = print_int s in begin
      dump dump_int (lstream ());
      assert (npeek 20 (lstream ()) = [1; 2; 3; 4]);
    end;
    
  dump dump_char (of_string "asdf");
  assert (npeek 20 (of_string "asdf") = ['a'; 's'; 'd'; 'f']);
  
  let o = open_out "/tmp/testdata" in begin
      output_string o "Test!";
      close_out o
    end;
  let i = open_in "/tmp/testdata" in begin
      dump dump_char (of_channel i);
      close_in i;
    end;
  let i = open_in "/tmp/testdata" in begin
      assert (npeek 20 (of_channel i) = ['T'; 'e'; 's'; 't'; '!']);
      close_in i;
    end;

  (* Stream should be the old implementation. *)
  let rec compare_streams (s1: 'a t) (s2: 'a Stream.t): bool =
    match (peek s1, Stream.peek s2) with
        (Some x, Some y) when x = y ->
          junk s1;
          Stream.junk s2;
          compare_streams s1 s2
      | (None, None) -> true
      | _ -> false
          (* Any file exceeding the size of the buffer is a good test here. *)
  in let myfilename = "stream_fixed.ml"
  in let i1 = open_in "stream_fixed.ml"
  in let s1 = of_channel i1
  in let i2 = open_in "stream_fixed.ml" 
  in let s2 = Stream.of_channel i2
  in let result = compare_streams s1 s2
  in begin
      close_in i1;
      close_in i2;
      assert result;
    end;
  assert (not (compare_streams (of_string "asdf") (Stream.of_string "fdsa")));

  let total = ref 0 in begin
      iter (function x -> total := x + !total) (of_list [1;2;3;4]);
      assert (10 = !total);
    end;
                         
  let s = of_list [1;2;3;4]
    in begin
      assert (1 == next s);
      assert (2 == next s);
      ignore (next s);
      ignore (next s);
      checkfailure (function _ -> next s);
    end;

  let s = of_list [1;2;3;4] in begin
      checkfailure (function _ -> empty s);
      assert (count s = 0);
      assert (next s = 1);
      assert (npeek 2 s = [2; 3]);
      assert (next s = 2);
      assert (count s = 2);
      assert (next s = 3);
      assert (peek s = Some 4);
      assert (count s = 3);
      junk s;
      assert (peek s = None);
      empty s;
    end;

  let s = iapp (of_list [3; 4]) (of_list [5; 6]) in begin
      assert (npeek 3 s = [3; 4; 5]);
      assert (count s = 0);
      assert (npeek 4 s = [3; 4; 5; 6]);
      junk s;
      junk s;
      assert (count s = 2);
      assert (npeek 4 s = [5; 6]);
      assert (npeek 12 (ising 3) = [3]);
    end;

  let i1 = open_in "/tmp/testdata" in
  let s1 = of_channel i1 in begin
      assert (npeek 4 (lapp (function _ -> (ising 'c'))
                         (slazy (function _ -> s1)))
              = ['c'; 'T'; 'e'; 's']);
      close_in i1;
      assert (npeek 4 (lsing (function _ -> 3)) = [3]);
      assert (npeek 4 sempty = []);
    end;

  (* The previous version raises Failure "illegal stream concatenation"
     for the next test. *)
  let i = open_in "/tmp/testdata" in begin
      assert (npeek 20 (iapp (of_string "as")
                          (icons 'z' (iapp (of_list ['d'; 'f'])
                                        (of_channel i)))) =
                ['a'; 's'; 'z'; 'd'; 'f'; 'T'; 'e'; 's'; 't'; '!']);
      close_in i;
    end;
  end
in big_test ();;

(* For timing runs:
   1. Make sure this file is named "stream_fixed.ml" and the interface is
      "stream_fixed.mli".  Your ocaml interpreter should have the old version
      of Stream, so we can compare the behavior of the two. 
   2. Put the following into timerun.ml.
   3. ocamlc stream_fixed.mli stream_fixed.ml timerun.ml
   4. time ./a.out.  (13.18 13.17 13.19 seconds with ocamlc; avg 13.18,
                      1.17 1.18 1.18 seconds with ocamlopt; avg 1.18.)
   5. Change Stream_fixed to Stream in timerun.ml.
   6. Give the same compilation command again.
   7. time ./a.out.  (14.61 14.62 14.64 seconds with ocamlc; avg 14.62,
                      1.11 1.11 1.10 seconds with ocamlopt.; avg 1.11.)
   So it's 10% faster when running interpreted, or 6% slower when running
   compiled.
 *)
module S = Stream_fixed;;
let myfilename = "stream_fixed.ml";;

let finally (body: unit -> 'a) (handler: unit -> unit): 'a =
  let result: 'a option ref = ref None
  in begin
      begin
        try begin
          result := Some (body ());
        end with e -> begin
          handler ();
          raise e;
        end;
      end;
      handler ();
      match !result with
          Some r -> r
        | _ -> failwith "Confused in finally";
    end;;

for i = 0 to 1000 do
   let i1 = open_in myfilename
   in finally
        (function _ ->
           let s1 = S.of_channel i1 in
             while None <> S.peek s1 do
               S.junk s1;
             done)
        (function _ -> close_in i1)
done;;

*)